catalog/lease: properly handle session ID changes

Previously, the lease manager did not properly handle cases when the
session ID changes. This could happen because of failover scenarios,
where a new session ID would be assigned. The lease manager would not
update the in memory or on disk state to pick up the new session ID, if
the descriptor version did not change.
This could lead to an infinite loop inside lease acquisition. To address
this, this patch will allow upserting leases with a new session ID and
the same version.

Fixes: #141567
Fixes: #141556
Fixes: #141555
Fixes: #141554
Fixes: #141553
Fixes: #141552
Fixes: #141549
Fixes: #141548
Fixes: #141547
Fixes: #141546
Fixes: #141545
Fixes: #141544
Fixes: #141543
Fixes: #141542
Fixes: #141541
Fixes: #141540
Fixes: #141539
Fixes: #141538
Fixes: #141484
Fixes: #141481
Fixes: #141480
Fixes :#141473
Fixes: #141473
Fixes: #141467
Fixes: #141685
Fixes: #141585
Fixes: #141566
Fixes: #141513
Fixes: #141479

Release note: None

Author: fqazi

pkg/sql/catalog/lease/BUILD.bazel

@@ -127,6 +127,7 @@ go_test(
         "//pkg/sql/sqlliveness",
         "//pkg/sql/sqlliveness/slbase",
         "//pkg/sql/sqlliveness/slprovider",
+        "//pkg/sql/sqlliveness/sqllivenesstestutils",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/stats",
         "//pkg/sql/types",

pkg/sql/catalog/lease/descriptor_state.go

@@ -125,7 +125,7 @@ func (t *descriptorState) upsertLeaseLocked(
 	desc catalog.Descriptor,
 	session sqlliveness.Session,
 	regionEnumPrefix []byte,
-) (createdDescriptorVersionState *descriptorVersionState, _ error) {
+) error {
 	if t.mu.maxVersionSeen < desc.GetVersion() {
 		t.mu.maxVersionSeen = desc.GetVersion()
 	}
@@ -136,10 +136,28 @@ func (t *descriptorState) upsertLeaseLocked(
 		}
 		descState := newDescriptorVersionState(t, desc, hlc.Timestamp{}, session, regionEnumPrefix, true /* isLease */)
 		t.mu.active.insert(descState)
-		return descState, nil
+		t.m.names.insert(ctx, descState)
+		return nil
+	}
+	// If the version already exists and the session ID matches nothing
+	// needs to be done.
+	if s.getSessionID() == session.ID() {
+		return nil
 	}
-	// If the version already exists, nothing needs to be done.
-	return nil, nil
+
+	// Otherwise, we need to update the existing lease to fix the session ID. The
+	// previously stored lease also needs to be deleted.
+	var existingLease storedLease
+	func() {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		existingLease = *s.mu.lease
+		s.mu.lease.sessionID = session.ID().UnsafeBytes()
+		s.mu.session = session
+	}()
+	// Delete the existing lease on behalf of the caller.
+	t.m.storage.release(ctx, t.m.stopper, &existingLease)
+	return nil
 }
 
 var _ redact.SafeFormatter = (*descriptorVersionState)(nil)

pkg/sql/catalog/lease/descriptor_version_state.go

@@ -108,6 +108,13 @@ func (s *descriptorVersionState) stringLocked() redact.RedactableString {
 	return redact.Sprintf("%d(%q,%s) ver=%d:%s, refcount=%d", s.GetID(), s.GetName(), redact.SafeString(sessionID), s.GetVersion(), s.mu.expiration, s.mu.refcount)
 }
 
+// getSessionID returns the current session ID from the lease.
+func (s *descriptorVersionState) getSessionID() sqlliveness.SessionID {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.mu.session.ID()
+}
+
 // hasExpired checks if the descriptor is too old to be used (by a txn
 // operating) at the given timestamp.
 func (s *descriptorVersionState) hasExpired(ctx context.Context, timestamp hlc.Timestamp) bool {

pkg/sql/catalog/lease/lease.go

@@ -853,15 +853,17 @@ func acquireNodeLease(
 			}
 			newest := m.findNewest(id)
 			var currentVersion descpb.DescriptorVersion
+			var currentSessionID sqlliveness.SessionID
 			if newest != nil {
 				currentVersion = newest.GetVersion()
+				currentSessionID = newest.getSessionID()
 			}
 			// A session will always be populated, since we use session based leasing.
 			session, err := m.storage.livenessProvider.Session(ctx)
 			if err != nil {
 				return false, errors.Wrapf(err, "lease acquisition was unable to resolve liveness session")
 			}
-			desc, regionPrefix, err := m.storage.acquire(ctx, session, id, currentVersion)
+			desc, regionPrefix, err := m.storage.acquire(ctx, session, id, currentVersion, currentSessionID)
 			if err != nil {
 				return nil, err
 			}
@@ -877,14 +879,10 @@ func acquireNodeLease(
 			t.mu.Lock()
 			t.mu.takenOffline = false
 			defer t.mu.Unlock()
-			var newDescVersionState *descriptorVersionState
-			newDescVersionState, err = t.upsertLeaseLocked(ctx, desc, session, regionPrefix)
+			err = t.upsertLeaseLocked(ctx, desc, session, regionPrefix)
 			if err != nil {
 				return nil, err
 			}
-			if newDescVersionState != nil {
-				m.names.insert(ctx, newDescVersionState)
-			}
 			return true, nil
 		})
 	if m.testingKnobs.LeaseStoreTestingKnobs.LeaseAcquireResultBlockEvent != nil {

pkg/sql/catalog/lease/lease_internal_test.go

@@ -29,12 +29,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/enum"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/sqllivenesstestutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
@@ -1606,3 +1608,119 @@ func TestLeaseCountDetailSessionBased(t *testing.T) {
 	require.Equal(t, 1, detail.numSQLInstances)
 	require.Equal(t, 0, detail.sampleSQLInstanceID)
 }
+
+// fakeSessionProvider session provider that only overloads the Session function
+// with a callback.
+type fakeSessionProvider struct {
+	syncutil.Mutex
+	sqlliveness.Provider
+	getSession func() *sqllivenesstestutils.FakeSession
+}
+
+var _ sqlliveness.Provider = &fakeSessionProvider{}
+
+// Session implements sqlliveness.Provider
+func (p *fakeSessionProvider) Session(ctx context.Context) (sqlliveness.Session, error) {
+	p.Lock()
+	defer p.Unlock()
+	if f := p.getSession(); f != nil {
+		return f, nil
+	}
+	return p.Provider.Session(ctx)
+}
+
+// TestLeaseManagerSessionIDChanges validates that the lease manager can acquire
+// and release leases properly even if the SessionID changes. This can happen
+// during a fail over scenario, where a new SessionID could be picked up. Which,
+// should cause us to reacquire leases.
+func TestLeaseManagerSessionIDChanges(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	srv, sqlDB, kvDB := serverutils.StartServer(
+		t, base.TestServerArgs{
+			// Avoid using tenants since async tenant migration steps can acquire
+			// leases on our user tables.
+			DefaultTestTenant: base.TestNeedsTightIntegrationBetweenAPIsAndTestingKnobs,
+		})
+	defer srv.Stopper().Stop(context.Background())
+	s := srv.ApplicationLayer()
+	leaseManager := s.LeaseManager().(*Manager)
+
+	runner := sqlutils.MakeSQLRunner(sqlDB)
+	runner.Exec(t, `SET CLUSTER SETTING sql.stats.automatic_collection.enabled = false`)
+	runner.Exec(t, `SET CLUSTER SETTING sql.catalog.descriptor_wait_for_initial_version.enabled = false`)
+
+	runner.Exec(t, `
+CREATE DATABASE t;
+CREATE TABLE t.test (k CHAR PRIMARY KEY, v CHAR);
+`)
+	tableDesc := desctestutils.TestingGetPublicTableDescriptor(kvDB, s.Codec(), "t", "test")
+
+	getLatestLeasedDesc := func() *descriptorVersionState {
+		state := leaseManager.findDescriptorState(tableDesc.GetID(), false)
+		state.mu.Lock()
+		defer state.mu.Unlock()
+		descState := state.mu.active.findNewest()
+		return descState
+	}
+
+	// Set up a fake session provider that will keep changing IDs and every session
+	// will instantly expire. This is like having a zero duration lease in the expiry
+	// model.
+	var nextSessionID atomic.Int64
+	var enableHook atomic.Bool
+
+	fs := fakeSessionProvider{
+		Provider: leaseManager.storage.livenessProvider,
+		getSession: func() *sqllivenesstestutils.FakeSession {
+			if !enableHook.Load() {
+				return nil
+			}
+			now := s.Clock().Now()
+			return &sqllivenesstestutils.FakeSession{
+				SessionID: sqlliveness.SessionID(fmt.Sprintf("session-%d", nextSessionID.Load())),
+				ExpTS:     now,
+				StartTS:   now,
+			}
+		},
+	}
+
+	// Replace the session provider which only returns expired leases.
+	leaseManager.storage.livenessProvider = &fs
+	defer func() {
+		// Restore the original provider so valid session IDs
+		// are assigned.
+		leaseManager.storage.livenessProvider = fs.Provider
+	}()
+
+	// Repeatedly lease the same descriptor, with the session ID continuously changing
+	// and each one always being expired. This validates that in fail over scenarios,
+	// nothing bad happens if the session is expired.
+	ctx := context.Background()
+	var previousSessionID sqlliveness.SessionID
+	var previousExpiry hlc.Timestamp
+	for count := 0; count < 10; count++ {
+		enableHook.Swap(true)
+		nextSessionID.Add(1)
+		now := s.Clock().Now()
+		desc, err := leaseManager.Acquire(ctx, now, tableDesc.GetID())
+		require.NoError(t, err)
+		// We expect a new session ID each time, and the descriptor
+		// to be expired.
+		newSessionID := getLatestLeasedDesc().getSessionID()
+		newExpiry := getLatestLeasedDesc().getExpiration(ctx)
+		require.NotEqualf(t, previousSessionID, newSessionID, "session ID should not match")
+		require.Truef(t, previousExpiry.Less(newExpiry), "session expiry should be later.")
+		// Disable the hook before the lease query.
+		enableHook.Swap(false)
+		// Sanity: Validate that system.lease only has a single lease with new
+		// id.
+		runner.CheckQueryResults(t, fmt.Sprintf("SELECT session_id FROM system.lease WHERE desc_id=%d", tableDesc.GetID()),
+			[][]string{{string(newSessionID.UnsafeBytes())}})
+		previousExpiry = newExpiry
+		previousSessionID = newSessionID
+		desc.Release(ctx)
+	}
+
+}

pkg/sql/catalog/lease/storage.go

@@ -120,14 +120,15 @@ func (s storage) crossValidateDuringRenewal() bool {
 // acquire a lease on the most recent version of a descriptor. The lease is tied
 // to the provided sqlliveness.Session. If a newer version (then lastVersion) of
 // the descriptor exists, this function will attempt to acquire a lease on it.
-// If no newer  version exists, it returns a nil descriptor. If the lease cannot
+// If no newer version exists, it returns a nil descriptor. If the lease cannot
 // be obtained because the descriptor is being dropped or is offline (currently
 // only applicable to tables), an inactiveTableError is returned.
 func (s storage) acquire(
 	ctx context.Context,
 	session sqlliveness.Session,
 	id descpb.ID,
 	lastVersion descpb.DescriptorVersion,
+	lastSessionID sqlliveness.SessionID,
 ) (desc catalog.Descriptor, prefix []byte, _ error) {
 	ctx = multitenant.WithTenantCostControlExemption(ctx)
 	prefix = s.getRegionPrefix()
@@ -151,7 +152,7 @@ func (s storage) acquire(
 		// written a value to the database, which we'd leak if we did not delete it.
 		// Note that the expiration is part of the primary key in the table, so we
 		// would not overwrite the old entry if we just were to do another insert.
-		//repeatIteration = desc != nil
+		// repeatIteration = desc != nil
 		if sessionID != nil && desc != nil {
 			if err := s.writer.deleteLease(ctx, txn, leaseFields{
 				regionPrefix: prefix,
@@ -168,10 +169,17 @@ func (s storage) acquire(
 		// any retryable error. If we run into an error then the delete
 		// above may need to be executed again.
 		latestDesc, err := s.mustGetDescriptorByID(ctx, txn, id)
-		// If the descriptor version hasn't changed, then nothing needs to be done.
-		if err != nil || latestDesc.GetVersion() == lastVersion {
+
+		if err != nil {
 			return err
 		}
+		// If the descriptor version hasn't changed, then no new lease to be
+		// inserted unless the session ID has changed on us. No descriptor will
+		// be set indicating to the caller that no new version exists.
+		if latestDesc.GetVersion() == lastVersion &&
+			lastSessionID == session.ID() {
+			return nil
+		}
 		desc = latestDesc
 		if err := catalog.FilterAddingDescriptor(desc); err != nil {
 			return err
@@ -234,6 +242,11 @@ func (s storage) acquire(
 		case err != nil:
 			return nil, nil, err
 		}
+		// If desc is nil then no new descriptor was available to be leased.
+		// i.e. the last version we leased is the latest and still held
+		if desc == nil {
+			return nil, nil, nil
+		}
 		log.VEventf(ctx, 2, "storage acquired lease %v", desc)
 		if s.testingKnobs.LeaseAcquiredEvent != nil {
 			s.testingKnobs.LeaseAcquiredEvent(desc, err)