Merge #153619

153619: catalog/lease: initial work for support timestamp locking leasing r=fqazi a=fqazi

Previously, when descriptors were leased it was possible to lease different versions of dependent descriptors, since the range feed used by the lease manager would make them instantly available. This meant that using leased descriptors for crdb_internal / pg_catalog functions with on going schema changes was not safe. To address this, this patch lays the initial ground work for locking leased descriptors to timestamps that the lease manager is aware of as consistent (i.e. all updates are received).

Informs: #153618
Release note: None

Co-authored-by: Faizan Qazi <faizan@cockroachlabs.com>

Author: craig[bot]

pkg/ccl/changefeedccl/cdcevent/rowfetcher_cache.go

@@ -74,7 +74,7 @@ func (f *dbTableDescFetcher) FetchTableDesc(
 ) (catalog.TableDescriptor, error) {
 	// Retrieve the target TableDescriptor from the lease manager. No caching
 	// is attempted because the lease manager does its own caching.
-	desc, err := f.leaseMgr.Acquire(ctx, ts, tableID)
+	desc, err := f.leaseMgr.Acquire(ctx, lease.TimestampToReadTimestamp(ts), tableID)
 	if err != nil {
 		// Manager can return all kinds of errors during chaos, but based on
 		// its usage, none of them should ever be terminal.

pkg/ccl/changefeedccl/schemafeed/schema_feed.go

@@ -58,7 +58,7 @@ func (e TableEvent) Timestamp() hlc.Timestamp {
 // leaseAcquirer is an interface containing the methods on *lease.Manager used
 // by the schema feed.
 type leaseAcquirer interface {
-	Acquire(ctx context.Context, timestamp hlc.Timestamp, id descpb.ID) (lease.LeasedDescriptor, error)
+	Acquire(ctx context.Context, timestamp lease.ReadTimestamp, id descpb.ID) (lease.LeasedDescriptor, error)
 	AcquireFreshestFromStore(ctx context.Context, id descpb.ID) error
 	// TODO(yang): Investigate whether the codec can be stored in the schema feed itself.
 	Codec() keys.SQLCodec
@@ -504,7 +504,7 @@ func (tf *schemaFeed) pauseOrResumePolling(ctx context.Context, atOrBefore hlc.T
 
 	if canPausePolling, err := tf.targets.EachTableIDWithBool(func(id descpb.ID) (bool, error) {
 		// Check if target table is schema-locked at the current frontier.
-		ld1, err := tf.leaseMgr.Acquire(ctx, frontier, id)
+		ld1, err := tf.leaseMgr.Acquire(ctx, lease.TimestampToReadTimestamp(frontier), id)
 		if err != nil {
 			return false, err
 		}
@@ -522,7 +522,7 @@ func (tf *schemaFeed) pauseOrResumePolling(ctx context.Context, atOrBefore hlc.T
 		}
 
 		// Check if target table remains at the same version at atOrBefore.
-		ld2, err := tf.leaseMgr.Acquire(ctx, atOrBefore, id)
+		ld2, err := tf.leaseMgr.Acquire(ctx, lease.TimestampToReadTimestamp(atOrBefore), id)
 		if err != nil {
 			return false, err
 		}

pkg/ccl/changefeedccl/schemafeed/schema_feed_test.go

@@ -323,14 +323,14 @@ type testLeaseAcquirer struct {
 }
 
 func (t *testLeaseAcquirer) Acquire(
-	ctx context.Context, timestamp hlc.Timestamp, id descpb.ID,
+	ctx context.Context, timestamp lease.ReadTimestamp, id descpb.ID,
 ) (lease.LeasedDescriptor, error) {
 	if id != t.id {
 		return nil, errors.Newf("unknown id: %d", id)
 	}
 
-	i, ok := slices.BinarySearchFunc(t.descs, timestamp, func(desc *testLeasedDescriptor, timestamp hlc.Timestamp) int {
-		return desc.timestamp.Compare(timestamp)
+	i, ok := slices.BinarySearchFunc(t.descs, timestamp, func(desc *testLeasedDescriptor, timestamp lease.ReadTimestamp) int {
+		return desc.timestamp.Compare(timestamp.GetTimestamp())
 	})
 	if ok {
 		return t.descs[i], nil

pkg/crosscluster/logical/lww_kv_processor.go

@@ -384,7 +384,7 @@ func (p *kvRowProcessor) getWriter(
 		}
 	}
 
-	l, err := p.cfg.LeaseManager.(*lease.Manager).Acquire(ctx, ts, id)
+	l, err := p.cfg.LeaseManager.(*lease.Manager).Acquire(ctx, lease.TimestampToReadTimestamp(ts), id)
 	if err != nil {
 		return nil, err
 	}

pkg/crosscluster/logical/tombstone_updater.go

@@ -159,7 +159,7 @@ func (tu *tombstoneUpdater) getDeleter(ctx context.Context, txn *kv.Txn) (row.De
 		tu.ReleaseLeases(ctx)
 
 		var err error
-		tu.leased.descriptor, err = tu.leaseMgr.Acquire(ctx, timestamp, tu.descID)
+		tu.leased.descriptor, err = tu.leaseMgr.Acquire(ctx, lease.TimestampToReadTimestamp(timestamp), tu.descID)
 		if err != nil {
 			return row.Deleter{}, err
 		}

pkg/sql/catalog/descs/collection.go

@@ -1373,7 +1373,7 @@ func (tc *Collection) GetIndexComment(
 // MaybeSetReplicationSafeTS modifies a txn to apply the replication safe timestamp,
 // if we are executing against a PCR reader catalog.
 func (tc *Collection) MaybeSetReplicationSafeTS(ctx context.Context, txn *kv.Txn) error {
-	now := txn.DB().Clock().Now()
+	now := tc.leased.lm.GetReadTimestamp(txn.DB().Clock().Now())
 	desc, err := tc.leased.lm.Acquire(ctx, now, keys.SystemDatabaseID)
 	if err != nil {
 		return err

pkg/sql/catalog/descs/leased_descriptors.go

@@ -23,14 +23,14 @@ import (
 type LeaseManager interface {
 	AcquireByName(
 		ctx context.Context,
-		timestamp hlc.Timestamp,
+		timestamp lease.ReadTimestamp,
 		parentID descpb.ID,
 		parentSchemaID descpb.ID,
 		name string,
 	) (lease.LeasedDescriptor, error)
 
 	Acquire(
-		ctx context.Context, timestamp hlc.Timestamp, id descpb.ID,
+		ctx context.Context, timestamp lease.ReadTimestamp, id descpb.ID,
 	) (lease.LeasedDescriptor, error)
 
 	IncGaugeAfterLeaseDuration(
@@ -40,10 +40,13 @@ type LeaseManager interface {
 	GetSafeReplicationTS() hlc.Timestamp
 
 	GetLeaseGeneration() int64
+
+	GetReadTimestamp(timestamp hlc.Timestamp) lease.ReadTimestamp
 }
 
 type deadlineHolder interface {
 	ReadTimestamp() hlc.Timestamp
+	ReadTimestampFixed() bool
 	UpdateDeadline(ctx context.Context, deadline hlc.Timestamp) error
 }
 
@@ -59,6 +62,10 @@ func (m maxTimestampBoundDeadlineHolder) ReadTimestamp() hlc.Timestamp {
 	return m.maxTimestampBound.Prev()
 }
 
+func (m maxTimestampBoundDeadlineHolder) ReadTimestampFixed() bool {
+	return true
+}
+
 // UpdateDeadline implements the deadlineHolder interface.
 func (m maxTimestampBoundDeadlineHolder) UpdateDeadline(
 	ctx context.Context, deadline hlc.Timestamp,
@@ -75,8 +82,10 @@ func makeLeasedDescriptors(lm LeaseManager) leasedDescriptors {
 // leasedDescriptors holds references to all the descriptors leased in the
 // transaction, and supports access by name and by ID.
 type leasedDescriptors struct {
-	lm    LeaseManager
-	cache nstree.NameMap
+	lm                LeaseManager
+	cache             nstree.NameMap
+	leaseTimestamp    lease.ReadTimestamp
+	leaseTimestampSet bool
 }
 
 // mismatchedExternalDataRowTimestamp is generated when the external row data timestamps
@@ -170,6 +179,22 @@ func (ld *leasedDescriptors) maybeAssertExternalRowDataTS(desc catalog.Descripto
 	})
 }
 
+// maybeInitReadTimestamp selects a read timestamp for the lease manager.
+func (ld *leasedDescriptors) maybeInitReadTimestamp(txn deadlineHolder) {
+	if ld.leaseTimestampSet {
+		return
+	}
+	readTimestamp := txn.ReadTimestamp()
+	ld.leaseTimestampSet = true
+	// Fixed timestamp queries will use descriptors at the user select timestamp.
+	if txn.ReadTimestampFixed() {
+		ld.leaseTimestamp = lease.TimestampToReadTimestamp(readTimestamp)
+		return
+	}
+	// Otherwise, get a safe read timestamp from the lease manager.
+	ld.leaseTimestamp = ld.lm.GetReadTimestamp(readTimestamp)
+}
+
 // getLeasedDescriptorByName return a leased descriptor valid for the
 // transaction, acquiring one if necessary. Due to a bug in lease acquisition
 // for dropped descriptors, the descriptor may have to be read from the store,
@@ -193,9 +218,8 @@ func (ld *leasedDescriptors) getByName(
 		desc = cached.(lease.LeasedDescriptor).Underlying()
 		return desc, false, nil
 	}
-
-	readTimestamp := txn.ReadTimestamp()
-	ldesc, err := ld.lm.AcquireByName(ctx, readTimestamp, parentID, parentSchemaID, name)
+	ld.maybeInitReadTimestamp(txn)
+	ldesc, err := ld.lm.AcquireByName(ctx, ld.leaseTimestamp, parentID, parentSchemaID, name)
 	const setTxnDeadline = true
 	return ld.getResult(ctx, txn, setTxnDeadline, ldesc, err)
 }
@@ -209,9 +233,8 @@ func (ld *leasedDescriptors) getByID(
 	if cached := ld.getCachedByID(ctx, id); cached != nil {
 		return cached, false, nil
 	}
-
-	readTimestamp := txn.ReadTimestamp()
-	desc, err := ld.lm.Acquire(ctx, readTimestamp, id)
+	ld.maybeInitReadTimestamp(txn)
+	desc, err := ld.lm.Acquire(ctx, ld.leaseTimestamp, id)
 	const setTxnDeadline = false
 	return ld.getResult(ctx, txn, setTxnDeadline, desc, err)
 }
@@ -363,6 +386,7 @@ func (ld *leasedDescriptors) releaseAll(ctx context.Context) {
 		return nil
 	})
 	ld.cache.Clear()
+	ld.leaseTimestampSet = false
 }
 
 func (ld *leasedDescriptors) release(ctx context.Context, descs []lease.IDVersion) {
@@ -371,6 +395,9 @@ func (ld *leasedDescriptors) release(ctx context.Context, descs []lease.IDVersio
 			removed.(lease.LeasedDescriptor).Release(ctx)
 		}
 	}
+	if ld.cache.Len() == 0 {
+		ld.leaseTimestampSet = false
+	}
 }
 
 func (ld *leasedDescriptors) numDescriptors() int {

pkg/sql/catalog/lease/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "kv_writer.go",
         "lease.go",
         "lease_test_utils.go",
+        "lease_timestamp.go",
         "name_cache.go",
         "storage.go",
         "testutils.go",

pkg/sql/catalog/lease/descriptor_state.go

@@ -82,7 +82,7 @@ type descriptorState struct {
 // It returns true if the descriptor returned is the known latest version
 // of the descriptor.
 func (t *descriptorState) findForTimestamp(
-	ctx context.Context, timestamp hlc.Timestamp,
+	ctx context.Context, timestamp ReadTimestamp,
 ) (*descriptorVersionState, bool, error) {
 	expensiveLogEnabled := log.ExpensiveLogEnabled(ctx, 2)
 	t.mu.Lock()
@@ -95,10 +95,12 @@ func (t *descriptorState) findForTimestamp(
 
 	// Walk back the versions to find one that is valid for the timestamp.
 	for i := len(t.mu.active.data) - 1; i >= 0; i-- {
-		// Check to see if the ModificationTime is valid.
-		if desc := t.mu.active.data[i]; desc.GetModificationTime().LessEq(timestamp) {
+		// Check to see if the ModificationTime is valid. If only the initial version
+		// of the descriptor is known, then read it at the base timestamp.
+		if desc := t.mu.active.data[i]; desc.GetModificationTime().LessEq(timestamp.GetTimestamp()) ||
+			(len(t.mu.active.data) == 1 && desc.GetModificationTime().LessEq(timestamp.GetBaseTimestamp())) {
 			latest := i+1 == len(t.mu.active.data)
-			if !desc.hasExpired(ctx, timestamp) {
+			if !desc.hasExpired(ctx, timestamp.GetBaseTimestamp()) {
 				// Existing valid descriptor version.
 				desc.incRefCount(ctx, expensiveLogEnabled)
 				return desc, latest, nil

pkg/sql/catalog/lease/lease.go

@@ -102,6 +102,12 @@ var WaitForInitialVersion = settings.RegisterBoolSetting(settings.ApplicationLev
 	"enables waiting for the initial version of a descriptor",
 	true)
 
+var LockedLeaseTimestamp = settings.RegisterBoolSetting(settings.ApplicationLevel,
+	"sql.catalog.descriptor_lease.use_locked_timestamps.enabled",
+	"guarantees transactional version consistency for descriptors used by the lease manager,"+
+		"descriptors used can be intentionally older to support this",
+	false)
+
 // WaitForNoVersion returns once there are no unexpired leases left
 // for any version of the descriptor.
 func (m *Manager) WaitForNoVersion(
@@ -1265,7 +1271,7 @@ func (m *Manager) purgeOldVersions(
 		// Acquire a refcount on the descriptor on the latest version to maintain an
 		// active lease, so that it doesn't get released when removeInactives()
 		// is called below. Release this lease after calling removeInactives().
-		desc, _, err = t.findForTimestamp(ctx, m.storage.clock.Now())
+		desc, _, err = t.findForTimestamp(ctx, TimestampToReadTimestamp(m.storage.clock.Now()))
 		if err == nil || !errors.Is(err, errRenewLease) {
 			break
 		}
@@ -1600,7 +1606,7 @@ func (m *Manager) SetRegionPrefix(val []byte) {
 // id and fails because the id has been dropped by the TRUNCATE.
 func (m *Manager) AcquireByName(
 	ctx context.Context,
-	timestamp hlc.Timestamp,
+	timestamp ReadTimestamp,
 	parentID descpb.ID,
 	parentSchemaID descpb.ID,
 	name string,
@@ -1622,9 +1628,9 @@ func (m *Manager) AcquireByName(
 		return desc, nil
 	}
 	// Check if we have cached an ID for this name.
-	descVersion, _ := m.names.get(ctx, parentID, parentSchemaID, name, timestamp)
+	descVersion, _ := m.names.get(ctx, parentID, parentSchemaID, name, timestamp.GetTimestamp())
 	if descVersion != nil {
-		if descVersion.GetModificationTime().LessEq(timestamp) {
+		if descVersion.GetModificationTime().LessEq(timestamp.GetTimestamp()) {
 			return validateDescriptorForReturn(descVersion)
 		}
 		// m.names.get() incremented the refcount, we decrement it to get a new
@@ -1643,7 +1649,7 @@ func (m *Manager) AcquireByName(
 	// lease with at least a bit of lifetime left in it. So, we do it the hard
 	// way: look in the database to resolve the name, then acquire a new lease.
 	var err error
-	id, err := m.resolveName(ctx, timestamp, parentID, parentSchemaID, name)
+	id, err := m.resolveName(ctx, timestamp.GetTimestamp(), parentID, parentSchemaID, name)
 	if err != nil {
 		return nil, err
 	}
@@ -1777,7 +1783,7 @@ type LeasedDescriptor interface {
 // can only return an older version of a descriptor if the latest version
 // can be leased; as it stands a dropped descriptor cannot be leased.
 func (m *Manager) Acquire(
-	ctx context.Context, timestamp hlc.Timestamp, id descpb.ID,
+	ctx context.Context, timestamp ReadTimestamp, id descpb.ID,
 ) (LeasedDescriptor, error) {
 	for {
 		if m.IsDraining() {
@@ -1802,7 +1808,7 @@ func (m *Manager) Acquire(
 
 		case errors.Is(err, errReadOlderVersion):
 			// Read old versions from the store. This can block while reading.
-			versions, errRead := m.readOlderVersionForTimestamp(ctx, id, timestamp)
+			versions, errRead := m.readOlderVersionForTimestamp(ctx, id, timestamp.GetTimestamp())
 			if errRead != nil {
 				return nil, errRead
 			}
@@ -2135,6 +2141,9 @@ func (m *Manager) watchForUpdates(ctx context.Context) {
 	}
 
 	handleCheckpoint := func(ctx context.Context, checkpoint *kvpb.RangeFeedCheckpoint) {
+		if m.testingKnobs.TestingOnRangeFeedCheckPoint != nil {
+			m.testingKnobs.TestingOnRangeFeedCheckPoint()
+		}
 		// Track checkpoints that occur from the rangefeed to make sure progress
 		// is always made.
 		m.mu.Lock()
@@ -2908,6 +2917,23 @@ func (m *Manager) deleteOrphanedLeasesWithSameInstanceID(
 		instanceID, releasedCount.Load(), totalLeases)
 }
 
+// GetReadTimestamp returns a locked timestamp to use for lease management.
+func (m *Manager) GetReadTimestamp(timestamp hlc.Timestamp) ReadTimestamp {
+	if LockedLeaseTimestamp.Get(&m.settings.SV) {
+		replicationTS := m.GetSafeReplicationTS()
+		if !replicationTS.IsEmpty() && replicationTS.Less(timestamp) {
+			return LeaseTimestamp{
+				ReadTimestamp:  timestamp,
+				LeaseTimestamp: replicationTS,
+			}
+		}
+	}
+	// Fallback to existing behavior with timestamps.
+	return LeaseTimestamp{
+		ReadTimestamp: timestamp,
+	}
+}
+
 // TestingGetBoundAccount returns the bound account used by the lease manager.
 func (m *Manager) TestingGetBoundAccount() *mon.ConcurrentBoundAccount {
 	return m.boundAccount