Merge #115568

115568: sqlliveness: adopt regionliveness when querying if a session is alive  r=fqazi a=fqazi

This PR will start query the regionliveness table as a part of the protocol to confirm if a given sqlliveness session IsAlive. As a side effect of this change the jobs subsystem and other infrastructure will be able take into account regionliveness. To allow this changed to happen the following commits are included in this PR

1. Update sqlliveness so that the heartbeat / TTL settings are moved into their out subpackage, avoiding circular dependencies with regionliveness
2. regionliveness is refactored to issue direct KV calls so that lower level code can utilize this infrastructure
3. The sqlliveness logic for IsAlive now takes into account information from the regionliveness table for both probing and detecting if a session should be dead.
4. Some of our tests for regionliveness did not configure the system database appropriately, so the region information was not correctly propagated to the sqlliveness table.

Informs: #115563

Co-authored-by: Faizan Qazi <[email protected]>

pkg/BUILD.bazel

@@ -2156,6 +2156,7 @@ GO_TARGETS = [
     "//pkg/sql/sqlitelogictest/tests/local-vec-off:local-vec-off_test",
     "//pkg/sql/sqlitelogictest/tests/local:local_test",
     "//pkg/sql/sqlitelogictest:sqlitelogictest",
+    "//pkg/sql/sqlliveness/slbase:slbase",
     "//pkg/sql/sqlliveness/slinstance:slinstance",
     "//pkg/sql/sqlliveness/slinstance:slinstance_test",
     "//pkg/sql/sqlliveness/slprovider:slprovider",

pkg/ccl/backupccl/BUILD.bazel

@@ -291,7 +291,7 @@ go_test(
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",
         "//pkg/sql/sqlclustersettings",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/stats",
         "//pkg/storage",
         "//pkg/testutils",

pkg/ccl/backupccl/backup_tenant_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	_ "github.com/cockroachdb/cockroach/pkg/sql/importer"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlclustersettings"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/jobutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -101,7 +101,7 @@ func TestBackupSharedProcessTenantNodeDown(t *testing.T) {
 	// for instance-based planning to recognize the downed node.
 	sv := &tenantApp.ClusterSettings().SV
 	padding := 10 * time.Second
-	timeout := slinstance.DefaultTTL.Get(sv) + slinstance.DefaultHeartBeat.Get(sv) + padding
+	timeout := slbase.DefaultTTL.Get(sv) + slbase.DefaultHeartBeat.Get(sv) + padding
 	testutils.SucceedsWithin(t, func() error {
 		_, err := tenantDB.Exec("BACKUP INTO 'nodelocal://1/worker-failure'")
 		return err

pkg/ccl/kvccl/kvtenantccl/upgradeccl/BUILD.bazel

@@ -25,7 +25,7 @@ go_test(
         "//pkg/spanconfig",
         "//pkg/sql/sem/eval",
         "//pkg/sql/sqlinstance/instancestorage",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",

pkg/ccl/kvccl/kvtenantccl/upgradeccl/tenant_upgrade_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/spanconfig"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlinstance/instancestorage"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -448,8 +448,8 @@ func TestTenantUpgradeFailure(t *testing.T) {
 		// Shorten the reclaim loop so that terminated SQL servers don't block
 		// the upgrade from succeeding.
 		instancestorage.ReclaimLoopInterval.Override(ctx, &settings.SV, 250*time.Millisecond)
-		slinstance.DefaultTTL.Override(ctx, &settings.SV, 15*time.Second)
-		slinstance.DefaultHeartBeat.Override(ctx, &settings.SV, 500*time.Millisecond)
+		slbase.DefaultTTL.Override(ctx, &settings.SV, 15*time.Second)
+		slbase.DefaultHeartBeat.Override(ctx, &settings.SV, 500*time.Millisecond)
 		tenantStopper := stop.NewStopper()
 		// Initialize the version to the minimum it could be.
 		require.NoError(t, clusterversion.Initialize(ctx, v0, &settings.SV))

pkg/ccl/kvccl/kvtenantccl/upgradeinterlockccl/BUILD.bazel

@@ -30,7 +30,7 @@ go_test(
         "//pkg/sql/catalog/lease",
         "//pkg/sql/sessiondatapb",
         "//pkg/sql/sqlinstance/instancestorage",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/stats",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",  # keep

pkg/ccl/kvccl/kvtenantccl/upgradeinterlockccl/local_test_util_test.go

@@ -27,7 +27,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlinstance/instancestorage"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -114,8 +114,8 @@ func runTest(t *testing.T, variant sharedtestutil.TestVariant, test sharedtestut
 			ttlOverride *= 10
 		}
 		heartbeatOverride := ttlOverride / 10
-		slinstance.DefaultTTL.Override(ctx, &s.SV, ttlOverride)
-		slinstance.DefaultHeartBeat.Override(ctx, &s.SV, heartbeatOverride)
+		slbase.DefaultTTL.Override(ctx, &s.SV, ttlOverride)
+		slbase.DefaultHeartBeat.Override(ctx, &s.SV, heartbeatOverride)
 	}
 
 	// Initialize the version to the MinSupportedVersion so that we can perform

pkg/ccl/multiregionccl/BUILD.bazel

@@ -86,7 +86,7 @@ go_test(
         "//pkg/sql/sem/tree",
         "//pkg/sql/sqlinstance/instancestorage",
         "//pkg/sql/sqlliveness",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/sqlliveness/slstorage",
         "//pkg/sql/sqltestutils",
         "//pkg/testutils",

pkg/ccl/multiregionccl/regionliveness_test.go

@@ -32,7 +32,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/regionliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/regions"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlinstance/instancestorage"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -65,7 +65,7 @@ func TestRegionLivenessProber(t *testing.T) {
 	makeSettings := func() *cluster.Settings {
 		cs := cluster.MakeTestingClusterSettings()
 		instancestorage.ReclaimLoopInterval.Override(ctx, &cs.SV, 150*time.Millisecond)
-		slinstance.DefaultTTL.Override(ctx, &cs.SV, 10*time.Second)
+		slbase.DefaultTTL.Override(ctx, &cs.SV, 10*time.Second)
 		regionliveness.RegionLivenessEnabled.Override(ctx, &cs.SV, true)
 		return cs
 	}
@@ -88,23 +88,19 @@ func TestRegionLivenessProber(t *testing.T) {
 	var tenants []serverutils.ApplicationLayerInterface
 	var tenantSQL []*gosql.DB
 	blockProbeQuery := atomic.Bool{}
+	defer regionliveness.TestingSetProbeLivenessTimeout(500*time.Millisecond,
+		func() {
+			// Timeout attempts to probe intentionally.
+			if blockProbeQuery.Swap(false) {
+				time.Sleep(2 * time.Second)
+			}
+		})()
 
 	for _, s := range testCluster.Servers {
 		tenantArgs := base.TestTenantArgs{
 			Settings: makeSettings(),
 			TenantID: id,
 			Locality: s.Locality(),
-			TestingKnobs: base.TestingKnobs{
-				SQLExecutor: &sql.ExecutorTestingKnobs{
-					BeforeExecute: func(ctx context.Context, stmt string, descriptors *descs.Collection) {
-						const probeQuery = "SELECT count(*) FROM system.sql_instances WHERE crdb_region = $1::system.crdb_internal_region"
-						if strings.Contains(stmt, probeQuery) && blockProbeQuery.Swap(false) {
-							// Timeout this query intentionally.
-							time.Sleep(testingRegionLivenessProbeTimeout + time.Second)
-						}
-					},
-				},
-			},
 		}
 		ts, tenantDB := serverutils.StartTenant(t, s, tenantArgs)
 		tenants = append(tenants, ts)
@@ -126,7 +122,7 @@ func TestRegionLivenessProber(t *testing.T) {
 	cachedRegionProvider, err = regions.NewCachedDatabaseRegions(ctx, tenants[0].DB(), tenants[0].LeaseManager().(*lease.Manager))
 	require.NoError(t, err)
 	idb := tenants[0].InternalDB().(isql.DB)
-	regionProber := regionliveness.NewLivenessProber(idb, cachedRegionProvider, tenants[0].ClusterSettings())
+	regionProber := regionliveness.NewLivenessProber(idb.KV(), tenants[0].Codec(), cachedRegionProvider, tenants[0].ClusterSettings())
 	// Validates the expected regions versus the region liveness set.
 	checkExpectedRegions := func(expectedRegions []string, regions regionliveness.LiveRegions) {
 		require.Equalf(t, len(regions), len(expectedRegions),
@@ -188,6 +184,24 @@ func TestRegionLivenessProber(t *testing.T) {
 			if _, ok := regions[expectedRegions[1]]; ok {
 				return errors.AssertionFailedf("removed region detected %s", expectedRegions[1])
 			}
+			// Similarly query the unavailable physcial regions
+			unavailablePhysicalRegions, err := regionProber.QueryUnavailablePhysicalRegions(ctx, txn)
+			if err != nil {
+				return err
+			}
+			if len(unavailablePhysicalRegions) != 1 {
+				return errors.AssertionFailedf("physical region was not marked as unavailable")
+			}
+			// Validate the physical region marked as unavailable is correct
+			regionTypeDesc := cachedRegionProvider.GetRegionEnumTypeDesc()
+			for i := 0; i < regionTypeDesc.NumEnumMembers(); i++ {
+				if regionTypeDesc.GetMemberLogicalRepresentation(i) != expectedRegions[1] {
+					continue
+				}
+				if !unavailablePhysicalRegions.ContainsPhysicalRepresentation(string(regionTypeDesc.GetMemberPhysicalRepresentation(i))) {
+					return errors.AssertionFailedf("incorrect physical region was marked as unavailable %v", unavailablePhysicalRegions)
+				}
+			}
 			return nil
 		})
 	})
@@ -211,7 +225,7 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 	makeSettings := func() *cluster.Settings {
 		cs := cluster.MakeTestingClusterSettings()
 		instancestorage.ReclaimLoopInterval.Override(ctx, &cs.SV, 150*time.Millisecond)
-		slinstance.DefaultTTL.Override(ctx, &cs.SV, 10*time.Second)
+		slbase.DefaultTTL.Override(ctx, &cs.SV, 10*time.Second)
 		regionliveness.RegionLivenessEnabled.Override(ctx, &cs.SV, true)
 		return cs
 	}
@@ -222,6 +236,18 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 		"us-west",
 	}
 	detectLeaseWait := atomic.Bool{}
+	targetCount := atomic.Int64{}
+	var tenants []serverutils.ApplicationLayerInterface
+	var tenantSQL []*gosql.DB
+	defer regionliveness.TestingSetProbeLivenessTimeout(1*time.Second, func() {
+		if !detectLeaseWait.Load() {
+			return
+		}
+		time.Sleep(time.Second * 2)
+		targetCount.Swap(0)
+		detectLeaseWait.Swap(false)
+	})()
+
 	testCluster, _, cleanup := multiregionccltestutils.TestingCreateMultiRegionClusterWithRegionList(t,
 		expectedRegions,
 		1,
@@ -231,12 +257,7 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 
 	id, err := roachpb.MakeTenantID(11)
 	require.NoError(t, err)
-
-	var tenants []serverutils.ApplicationLayerInterface
-	var tenantSQL []*gosql.DB
-	targetCount := atomic.Int64{}
-
-	for _, s := range testCluster.Servers {
+	for i, s := range testCluster.Servers {
 		tenantArgs := base.TestTenantArgs{
 			Settings: makeSettings(),
 			TenantID: id,
@@ -247,18 +268,13 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 						if !detectLeaseWait.Load() {
 							return
 						}
-						const probeQuery = "SELECT count(*) FROM system.sql_instances WHERE crdb_region = $1::system.public.crdb_internal_region"
 						const leaseQuery = "SELECT count(1) FROM system.public.lease AS OF SYSTEM TIME"
 						// Fail intentionally, when we go to probe the first region.
 						if strings.Contains(stmt, leaseQuery) {
 							if targetCount.Add(1) != 1 {
 								return
 							}
-							time.Sleep(testingRegionLivenessProbeTimeout + time.Second)
-						} else if strings.Contains(stmt, probeQuery) {
-							time.Sleep(testingRegionLivenessProbeTimeout + time.Second)
-							targetCount.Swap(0)
-							detectLeaseWait.Swap(false)
+							time.Sleep(time.Second * 2)
 						}
 					},
 				},
@@ -267,16 +283,21 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 		tenant, tenantDB := serverutils.StartTenant(t, s, tenantArgs)
 		tenantSQL = append(tenantSQL, tenantDB)
 		tenants = append(tenants, tenant)
+		// Before the other tenants are added we need to configure the system database,
+		// otherwise they will come up in a non multi-region mode and not all subsystems
+		// will be aware (i.e. session ID and SQL instance will not be MR aware).
+		if i == 0 {
+			// Convert into a multi-region DB.
+			_, err = tenantSQL[0].Exec(fmt.Sprintf("ALTER DATABASE system SET PRIMARY REGION '%s'", testCluster.Servers[0].Locality().Tiers[0].Value))
+			require.NoError(t, err)
+			for i := 1; i < len(expectedRegions); i++ {
+				_, err = tenantSQL[0].Exec(fmt.Sprintf("ALTER DATABASE system ADD REGION '%s'", expectedRegions[i]))
+				require.NoError(t, err)
+			}
+			_, err = tenantSQL[0].Exec("ALTER DATABASE system SURVIVE ZONE FAILURE")
+			require.NoError(t, err)
+		}
 	}
-	// Convert into a multi-region DB.
-	_, err = tenantSQL[0].Exec(fmt.Sprintf("ALTER DATABASE system SET PRIMARY REGION '%s'", testCluster.Servers[0].Locality().Tiers[0].Value))
-	require.NoError(t, err)
-	for i := 1; i < len(expectedRegions); i++ {
-		_, err = tenantSQL[0].Exec(fmt.Sprintf("ALTER DATABASE system ADD REGION '%s'", expectedRegions[i]))
-		require.NoError(t, err)
-	}
-	_, err = tenantSQL[0].Exec("ALTER DATABASE system SURVIVE ZONE FAILURE")
-	require.NoError(t, err)
 	// Override the table timeout probe for testing.
 	for _, ts := range tenants {
 		regionliveness.RegionLivenessProbeTimeout.Override(ctx, &ts.ClusterSettings().SV, testingRegionLivenessProbeTimeout)
@@ -296,7 +317,7 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 	detectLeaseWait.Swap(true)
 	_, err = tenantSQL[1].Exec("ALTER TABLE t1 ADD COLUMN i INT")
 	require.ErrorContainsf(t, err, "count-lease timed out reading from a region", "failed to timeout")
-	// Keep an active lease on node 1, but it will be seen as ignored eventually
+	// Keep an active lease on node 0, but it will be seen as ignored eventually
 	// because the region will start to get quarantined.
 	tx, err := tenantSQL[0].Begin()
 	require.NoError(t, err)
@@ -316,7 +337,7 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 
 	require.NoError(t, tx.Rollback())
 
-	// Validate we can have a "droped" region and the query won't fail.
+	// Validate we can have a "dropped" region and the query won't fail.
 	lm := tenants[0].LeaseManager().(*lease.Manager)
 	cachedDatabaseRegions, err := regions.NewCachedDatabaseRegions(ctx, tenants[0].DB(), lm)
 	require.NoError(t, err)
@@ -327,7 +348,7 @@ func TestRegionLivenessProberForLeases(t *testing.T) {
 			PhysicalRepresentation: nil,
 			LogicalRepresentation:  "FakeRegion",
 		})
-		return mut
+		return builder.BuildExistingMutableType()
 	})
 	require.NoError(t, lm.WaitForNoVersion(ctx, descpb.ID(tableID), cachedDatabaseRegions, retry.Options{}))
 }

pkg/ccl/multitenantccl/tenantcostclient/BUILD.bazel

@@ -86,7 +86,7 @@ go_test(
         "//pkg/sql/execinfra",
         "//pkg/sql/sem/eval",
         "//pkg/sql/sqlliveness",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/stats",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",

pkg/ccl/multitenantccl/tenantcostclient/tenant_side_test.go

@@ -48,7 +48,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/distsql"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
@@ -1024,7 +1024,7 @@ func TestSQLLivenessExemption(t *testing.T) {
 	st := cluster.MakeTestingClusterSettings()
 	// Make the tenant heartbeat like crazy.
 	ctx := context.Background()
-	slinstance.DefaultHeartBeat.Override(ctx, &st.SV, 10*time.Millisecond)
+	slbase.DefaultHeartBeat.Override(ctx, &st.SV, 10*time.Millisecond)
 
 	_, tenantDB := serverutils.StartTenant(t, hostServer, base.TestTenantArgs{
 		TenantID: tenantID,

pkg/ccl/testccl/sqlccl/BUILD.bazel

@@ -44,7 +44,7 @@ go_test(
         "//pkg/sql/isql",
         "//pkg/sql/lexbase",
         "//pkg/sql/sessiondatapb",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/tests",
         "//pkg/testutils",

pkg/ccl/testccl/sqlccl/temp_table_clean_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -104,8 +104,8 @@ func TestTenantTempTableCleanup(t *testing.T) {
 		sql.TempObjectWaitInterval.Override(ctx, &st.SV, time.Second*0)
 		// Set up sessions to expire within 5 seconds of a
 		// nodes death.
-		slinstance.DefaultTTL.Override(ctx, &st.SV, 5*time.Second)
-		slinstance.DefaultHeartBeat.Override(ctx, &st.SV, time.Second)
+		slbase.DefaultTTL.Override(ctx, &st.SV, 5*time.Second)
+		slbase.DefaultHeartBeat.Override(ctx, &st.SV, time.Second)
 		return st
 	}
 

pkg/kv/kvserver/gc/BUILD.bazel

@@ -64,7 +64,7 @@ go_test(
         "//pkg/server",
         "//pkg/server/serverpb",
         "//pkg/settings/cluster",
-        "//pkg/sql/sqlliveness/slinstance",
+        "//pkg/sql/sqlliveness/slbase",
         "//pkg/storage",
         "//pkg/storage/enginepb",
         "//pkg/testutils",

pkg/kv/kvserver/gc/gc_int_test.go

@@ -27,7 +27,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slinstance"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slbase"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -86,7 +86,7 @@ func TestEndToEndGC(t *testing.T) {
 				settings := cluster.MakeTestingClusterSettings()
 				// Push the TTL up to 60 hours since we emulate a 50 hours
 				// clock jump below.
-				slinstance.DefaultTTL.Override(ctx, &settings.SV, 60*time.Hour)
+				slbase.DefaultTTL.Override(ctx, &settings.SV, 60*time.Hour)
 
 				manualClock := hlc.NewHybridManualClock()
 				s, appSqlDb, appKvDb := serverutils.StartServer(t, base.TestServerArgs{