catalog/lease: retain descriptor versions for old timestamps

Previously, when we added support locked timestamps leasing we would not
hold descriptors that previous versions that older timestamps may need.
This could lead to use needing to reading old versions from KV more
frequently as well as limiting when transactions could commit in this
mode. To address this, this patch allows the lease manager to hold older
versions while the timestamps is in use.

Fixes: #153618

Release note: None

Author: fqazi

pkg/sql/catalog/descs/collection.go

@@ -1482,7 +1482,8 @@ 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 := tc.leased.lm.GetReadTimestamp(txn.DB().Clock().Now())
+	now := tc.leased.lm.GetReadTimestamp(ctx, txn.DB().Clock().Now())
+	defer now.Release(ctx)
 	desc, err := tc.leased.lm.Acquire(ctx, now, keys.SystemDatabaseID)
 	if err != nil {
 		return err

pkg/sql/catalog/descs/leased_descriptors.go

@@ -41,7 +41,7 @@ type LeaseManager interface {
 
 	GetLeaseGeneration() int64
 
-	GetReadTimestamp(timestamp hlc.Timestamp) lease.ReadTimestamp
+	GetReadTimestamp(ctx context.Context, timestamp hlc.Timestamp) lease.ReadTimestamp
 }
 
 type deadlineHolder interface {
@@ -180,24 +180,27 @@ func (ld *leasedDescriptors) maybeAssertExternalRowDataTS(desc catalog.Descripto
 }
 
 // maybeInitReadTimestamp selects a read timestamp for the lease manager.
-func (ld *leasedDescriptors) maybeInitReadTimestamp(txn deadlineHolder) {
+func (ld *leasedDescriptors) maybeInitReadTimestamp(ctx context.Context, txn deadlineHolder) {
 	// Refresh the leased timestamp if the read timestamp has changed on us
 	// or if it hasn't been populated yet.
-	// TODO (fqazi): For locked read timestamps inside leasing,
-	// we will need to have extra logic to ensure this is safe.
 	if ld.leaseTimestampSet &&
 		ld.leaseTimestamp.GetBaseTimestamp() == txn.ReadTimestamp() {
 		return
+	} else if ld.leaseTimestampSet {
+		// Timestamp is already set so release it before picking a new one.
+		// TODO (fqazi): Add handling for moving timestamps forward.
+		ld.leaseTimestamp.Release(ctx)
+		ld.leaseTimestampSet = false
 	}
 	readTimestamp := txn.ReadTimestamp()
 	ld.leaseTimestampSet = true
-	// Fixed timestamp queries will use descriptors at the user select timestamp.
+	// Fixed timestamp queries will use descriptors at the user-selected 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)
+	ld.leaseTimestamp = ld.lm.GetReadTimestamp(ctx, readTimestamp)
 }
 
 // getLeasedDescriptorByName return a leased descriptor valid for the
@@ -223,7 +226,7 @@ func (ld *leasedDescriptors) getByName(
 		desc = cached.(lease.LeasedDescriptor).Underlying()
 		return desc, false, nil
 	}
-	ld.maybeInitReadTimestamp(txn)
+	ld.maybeInitReadTimestamp(ctx, txn)
 	ldesc, err := ld.lm.AcquireByName(ctx, ld.leaseTimestamp, parentID, parentSchemaID, name)
 	const setTxnDeadline = true
 	return ld.getResult(ctx, txn, setTxnDeadline, ldesc, err)
@@ -238,7 +241,7 @@ func (ld *leasedDescriptors) getByID(
 	if cached := ld.getCachedByID(ctx, id); cached != nil {
 		return cached, false, nil
 	}
-	ld.maybeInitReadTimestamp(txn)
+	ld.maybeInitReadTimestamp(ctx, txn)
 	desc, err := ld.lm.Acquire(ctx, ld.leaseTimestamp, id)
 	const setTxnDeadline = false
 	return ld.getResult(ctx, txn, setTxnDeadline, desc, err)
@@ -286,7 +289,7 @@ func (ld *leasedDescriptors) getResult(
 	expiration := ldesc.Expiration(ctx)
 	readTimestamp := txn.ReadTimestamp()
 	if expiration.LessEq(txn.ReadTimestamp()) {
-		return nil, false, errors.AssertionFailedf("bad descriptor for id=%d readTimestamp=%s, expiration=%s", ldesc.GetID(), readTimestamp, expiration)
+		return nil, false, errors.AssertionFailedf("bad descriptor for id=%d readTimestamp=%s, expiration=%s (leaseTimestamp: %s)", ldesc.GetID(), readTimestamp, expiration, ld.leaseTimestamp.GetTimestamp())
 	}
 
 	ld.cache.Upsert(ldesc, ldesc.Underlying().SkipNamespace())
@@ -391,6 +394,9 @@ func (ld *leasedDescriptors) releaseAll(ctx context.Context) {
 		return nil
 	})
 	ld.cache.Clear()
+	if ld.leaseTimestampSet {
+		ld.leaseTimestamp.Release(ctx)
+	}
 	ld.leaseTimestampSet = false
 }
 
@@ -401,6 +407,9 @@ func (ld *leasedDescriptors) release(ctx context.Context, descs []lease.IDVersio
 		}
 	}
 	if ld.cache.Len() == 0 {
+		if ld.leaseTimestampSet {
+			ld.leaseTimestamp.Release(ctx)
+		}
 		ld.leaseTimestampSet = false
 	}
 }

pkg/sql/catalog/lease/BUILD.bazel

@@ -124,6 +124,7 @@ go_test(
         "//pkg/sql/schemachanger/scexec",
         "//pkg/sql/schemachanger/scop",
         "//pkg/sql/schemachanger/scplan",
+        "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",
         "//pkg/sql/sqlinstance/instancestorage",

pkg/sql/catalog/lease/descriptor_set.go

@@ -91,6 +91,13 @@ func (l *descriptorSet) findNewest() *descriptorVersionState {
 	return l.data[len(l.data)-1]
 }
 
+func (l *descriptorSet) findPrevious() *descriptorVersionState {
+	if len(l.data) < 2 {
+		return nil
+	}
+	return l.data[len(l.data)-2]
+}
+
 func (l *descriptorSet) findPreviousToExpire(dropped bool) *descriptorVersionState {
 	// If there are no versions, then no previous version exists.
 	if len(l.data) == 0 {

pkg/sql/catalog/lease/descriptor_state.go

@@ -92,18 +92,36 @@ func (t *descriptorState) findForTimestamp(
 	if len(t.mu.active.data) == 0 {
 		return nil, false, errRenewLease
 	}
+	return t.findForTimestampImpl(ctx, timestamp.GetTimestamp(), timestamp.GetBaseTimestamp(), expensiveLogEnabled)
+}
 
+// findForTimestampImpl searches for a descriptor that is valid for both the lease
+// timestamp and read timestamp. If no such descriptor exists, then this function
+// will return a descriptor satisfying the read timestamp (when using locked leases
+// this means we will allow mixing versions on first reads).
+func (t *descriptorState) findForTimestampImpl(
+	ctx context.Context,
+	leaseTimestamp hlc.Timestamp,
+	readTimestamp hlc.Timestamp,
+	expensiveLogEnabled bool,
+) (*descriptorVersionState, bool, error) {
 	// 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 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())) {
+		if desc := t.mu.active.data[i]; desc.GetModificationTime().LessEq(leaseTimestamp) {
 			latest := i+1 == len(t.mu.active.data)
-			if !desc.hasExpired(ctx, timestamp.GetBaseTimestamp()) {
+			if !desc.hasExpired(ctx, readTimestamp) {
 				// Existing valid descriptor version.
 				desc.incRefCount(ctx, expensiveLogEnabled)
 				return desc, latest, nil
+			} else if !latest && readTimestamp != leaseTimestamp {
+				// The lease timestamp is not compatible with the read timestamp, since
+				// the descriptor returned will be expired. This means we are seeing the
+				// first read of this descriptor, since the prior version was not locked.
+				// In this scenario, we will allow this transaction to mix versions as a
+				// last resort.
+				return t.findForTimestampImpl(ctx, readTimestamp, readTimestamp, expensiveLogEnabled)
 			}
 
 			if latest {