sql: use LeaseManager of ApplicationLayer in tests

This commit audits all tests in `pkg/sql` directory to ensure the
LeaseManager of the application layer is used. Note that tests
still might need further modifications.

Release note: None

Author: yuzefovich

pkg/sql/catalog/lease/lease_internal_test.go

@@ -1786,6 +1786,7 @@ func TestLeaseManagerLockedTimestampBasic(t *testing.T) {
 			Settings: st,
 		})
 	defer srv.Stopper().Stop(context.Background())
+	s := srv.ApplicationLayer()
 
 	r := sqlutils.MakeSQLRunner(db)
 
@@ -1802,7 +1803,7 @@ func TestLeaseManagerLockedTimestampBasic(t *testing.T) {
 	}()
 	var id int
 	r.QueryRow(t, "SELECT 't1'::REGCLASS::OID;").Scan(&id)
-	lm := srv.LeaseManager().(*Manager)
+	lm := s.LeaseManager().(*Manager)
 
 	// Note: We only need to hold old leases because the logic for doing this
 	// automatically is not merged yet. Once it is merged, the release of them
@@ -1817,7 +1818,7 @@ func TestLeaseManagerLockedTimestampBasic(t *testing.T) {
 	defer releaseHeldDescriptors()
 
 	getDescriptorVersion := func() descpb.DescriptorVersion {
-		ts := lm.GetReadTimestamp(srv.Clock().Now())
+		ts := lm.GetReadTimestamp(s.Clock().Now())
 		state := lm.findDescriptorState(descpb.ID(id), false)
 		require.NotNilf(t, state, "the descriptor was not leased yet")
 		ld, _, err := state.findForTimestamp(ctx, ts)
@@ -1869,12 +1870,12 @@ func TestLeaseManagerLockedTimestampCluster(t *testing.T) {
 	tc := serverutils.StartCluster(
 		t, 3, base.TestClusterArgs{
 			ServerArgs: base.TestServerArgs{
-				Settings:          st,
+				Settings: st,
+				// Avoid using tenants since async tenant migration steps can acquire
+				// leases on our user tables.
 				DefaultTestTenant: base.TestNeedsTightIntegrationBetweenAPIsAndTestingKnobs,
 			},
 			ServerArgsPerNode: map[int]base.TestServerArgs{2: {
-				// Avoid using tenants since async tenant migration steps can acquire
-				// leases on our user tables.
 				Knobs: base.TestingKnobs{
 					SQLLeaseManager: &ManagerTestingKnobs{
 						TestingDescriptorRefreshedEvent: func(descriptor *descpb.Descriptor) {
@@ -1907,8 +1908,6 @@ func TestLeaseManagerLockedTimestampCluster(t *testing.T) {
 		_, err := node1Conn.Exec("ALTER TABLE d1.public.t1 ADD COLUMN n2 int DEFAULT 364")
 		return err
 	})
-	go func() {
-	}()
 	lm := tc.Server(2).LeaseManager().(*Manager)
 	assertDescriptorsCount := func(expectedCount int) {
 		state := lm.findDescriptorState(descpb.ID(id), false)

pkg/sql/catalog/lease/lease_test.go

@@ -1828,12 +1828,13 @@ func TestDeleteOrphanedLeasesBySession(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	ctx := context.Background()
-	s, conn, _ := serverutils.StartServer(t, base.TestServerArgs{
+	srv, conn, _ := serverutils.StartServer(t, base.TestServerArgs{
 		Locality: roachpb.Locality{
 			Tiers: []roachpb.Tier{{Key: "region", Value: "us-east1"}},
 		},
 	})
-	defer s.Stop(ctx)
+	defer srv.Stop(ctx)
+	s := srv.ApplicationLayer()
 	idb := s.InternalDB().(*sql.InternalDB)
 	ie := idb.Executor()
 	// Validate only one session exists.
@@ -2841,11 +2842,12 @@ func TestLeaseTxnDeadlineExtensionWithSession(t *testing.T) {
 		},
 	}
 
-	s, sqlDB, _ := serverutils.StartServer(t, params)
+	srv, sqlDB, _ := serverutils.StartServer(t, params)
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	lm := s.LeaseManager().(*lease.Manager)
 	leasesWaitingToExpireGauge := lm.TestingSessionBasedLeasesWaitingToExpireGauge()
 	leasesExpiredGauge := lm.TestingSessionBasedLeasesExpiredGauge()
-	defer s.Stopper().Stop(ctx)
 	// Setup tables for the test.
 	_, err := sqlDB.Exec(`
 CREATE TABLE t1(val int);
@@ -3501,6 +3503,7 @@ func TestLongLeaseWaitMetrics(t *testing.T) {
 		},
 	})
 	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	runner := sqlutils.MakeSQLRunner(sqlDB)
 	runner.Exec(t, "CREATE TABLE t1(n int)")
 	descIDRow := runner.QueryRow(t, "SELECT 't1'::REGCLASS::INT")
@@ -3509,7 +3512,7 @@ func TestLongLeaseWaitMetrics(t *testing.T) {
 	grp := ctxgroup.WithContext(ctx)
 
 	startWaiters := make(chan struct{})
-	cachedDatabaseRegions, err := regions.NewCachedDatabaseRegions(ctx, srv.DB(), srv.LeaseManager().(*lease.Manager))
+	cachedDatabaseRegions, err := regions.NewCachedDatabaseRegions(ctx, s.DB(), s.LeaseManager().(*lease.Manager))
 	require.NoError(t, err)
 
 	grp.GoCtx(func(ctx context.Context) error {
@@ -3529,13 +3532,13 @@ func TestLongLeaseWaitMetrics(t *testing.T) {
 		r := retry.StartWithCtx(ctx, retry.Options{})
 		// Wait until long waits are detected in our metrics.
 		for r.Next() {
-			if srv.MustGetSQLCounter("sql.leases.long_wait_for_no_version") == 0 {
+			if s.MustGetSQLCounter("sql.leases.long_wait_for_no_version") == 0 {
 				continue
 			}
-			if srv.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant") == 0 {
+			if s.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant") == 0 {
 				continue
 			}
-			if srv.MustGetSQLCounter("sql.leases.long_wait_for_one_version") == 0 {
+			if s.MustGetSQLCounter("sql.leases.long_wait_for_one_version") == 0 {
 				continue
 			}
 			break
@@ -3563,11 +3566,11 @@ func TestLongLeaseWaitMetrics(t *testing.T) {
 		r := retry.StartWithCtx(ctx, retry.Options{})
 		for r.Next() {
 			// Wait for the two versions to exist.
-			if srv.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant") == 0 {
+			if s.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant") == 0 {
 				continue
 			}
 			// Wait for there to be a single version.
-			lm := srv.ApplicationLayer().LeaseManager().(*lease.Manager)
+			lm := s.LeaseManager().(*lease.Manager)
 			_, err := lm.WaitForOneVersion(ctx, descpb.ID(descID), cachedDatabaseRegions, retry.Options{})
 			return err
 		}
@@ -3577,16 +3580,16 @@ func TestLongLeaseWaitMetrics(t *testing.T) {
 	// Waits for no version of the descriptor to exist.
 	grp.GoCtx(func(ctx context.Context) error {
 		<-startWaiters
-		lm := srv.ApplicationLayer().LeaseManager().(*lease.Manager)
+		lm := s.LeaseManager().(*lease.Manager)
 		return lm.WaitForNoVersion(ctx, descpb.ID(descID), cachedDatabaseRegions, retry.Options{})
 	})
 
 	require.NoError(t, grp.Wait())
 
 	// Validate the metrics are 0 again.
-	require.Equal(t, int64(0), srv.MustGetSQLCounter("sql.leases.long_wait_for_no_version"))
-	require.Equal(t, int64(0), srv.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant"))
-	require.Equal(t, int64(0), srv.MustGetSQLCounter("sql.leases.long_wait_for_one_version"))
+	require.Equal(t, int64(0), s.MustGetSQLCounter("sql.leases.long_wait_for_no_version"))
+	require.Equal(t, int64(0), s.MustGetSQLCounter("sql.leases.long_wait_for_two_version_invariant"))
+	require.Equal(t, int64(0), s.MustGetSQLCounter("sql.leases.long_wait_for_one_version"))
 }
 
 // TestWaitForInitialVersionConcurrent this test is a basic sanity test that
@@ -3672,8 +3675,9 @@ func BenchmarkAcquireLeaseConcurrent(b *testing.B) {
 	defer log.Scope(b).Close(b)
 
 	ctx := context.Background()
-	s, sqlDB, _ := serverutils.StartServer(b, base.TestServerArgs{})
-	defer s.Stopper().Stop(ctx)
+	srv, sqlDB, _ := serverutils.StartServer(b, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 
 	runner := sqlutils.MakeSQLRunner(sqlDB)
 
@@ -3710,8 +3714,9 @@ func TestLeaseManagerIsMemoryMonitored(t *testing.T) {
 	skip.UnderDuress(t)
 
 	ctx := context.Background()
-	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
-	defer s.Stopper().Stop(ctx)
+	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	lm := s.LeaseManager().(*lease.Manager)
 	startBytes := lm.TestingGetBoundAccount().Used()
 	lastBytes := startBytes

pkg/sql/indexbackfiller_test.go

@@ -482,8 +482,8 @@ INSERT INTO foo VALUES (1), (10), (100);
 
 		// Run the testCase's setupDesc function to prepare an index backfill
 		// mutation. Also, create an associated job and set it up to be blocked.
-		lm := s.LeaseManager().(*lease.Manager)
 		tt := s.ApplicationLayer()
+		lm := tt.LeaseManager().(*lease.Manager)
 		codec := tt.Codec()
 		settings := tt.ClusterSettings()
 		execCfg := tt.ExecutorConfig().(sql.ExecutorConfig)

pkg/sql/schema_changer_test.go

@@ -1543,13 +1543,14 @@ func TestSchemaChangeRetryOnVersionChange(t *testing.T) {
 		JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 	}
 
-	s, sqlDB, kvDB := serverutils.StartServer(t, params)
-	defer s.Stopper().Stop(context.Background())
+	srv, sqlDB, kvDB := serverutils.StartServer(t, params)
+	defer srv.Stopper().Stop(context.Background())
 	defer func() {
 		t.Log("unblocking GC")
 		close(unblockGC)
 	}()
-	codec := s.ApplicationLayer().Codec()
+	s := srv.ApplicationLayer()
+	codec := s.Codec()
 
 	if _, err := sqlDB.Exec(`
 SET create_table_with_schema_locked=false;

pkg/sql/schemachanger/scexec/executor_external_test.go

@@ -92,7 +92,7 @@ func setupTestInfra(t testing.TB) (_ *testInfra, cleanup func(context.Context))
 		nodeID:   s.NodeID(),
 		settings: tt.ClusterSettings(),
 		db:       tt.ExecutorConfig().(sql.ExecutorConfig).InternalDB,
-		lm:       s.LeaseManager().(*lease.Manager),
+		lm:       tt.LeaseManager().(*lease.Manager),
 		cf:       tt.ExecutorConfig().(sql.ExecutorConfig).CollectionFactory,
 		tsql:     sqlutils.MakeSQLRunner(db),
 	}, s.Stopper().Stop