sql: force PCR reader catalogs to use a fixed timestamp

Previously, the original logic for advancing the PCR reader catalog
timestamp ran into the risk of mixing different timestamps, since the
lease manager relies on range feeds. So, it was possible for a single
txn to observe some newer and older descriptors. To address this, this
patch does the following:

1) It introduces an assertion in the leased descriptors to ensure that
   all external row data tables have the same timestamp.
2) The connection executor is modified for PCR reader catalog to always
   use AOST queries, where the AOST timestamp is determine based on the
   range feeds check point. This ensures that all data up to the
   timestamp synced up. We detect PCR reader catalog by checking the
   the ReplicatedPCRVersion on the systme database, which will be set
   during bootstrap.

Fixes: #130812

Release note: None

Author: fqazi

pkg/sql/catalog/descs/leased_descriptors.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/nstree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/iterutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 )
@@ -79,6 +80,93 @@ type leasedDescriptors struct {
 	cache nstree.NameMap
 }
 
+// mismatchedExternalDataRowTimestamp is generated when the external row data timestamps
+// within a descriptor do not match.
+type mismatchedExternalDataRowTimestamp struct {
+	newDescName      string
+	newDescID        descpb.ID
+	newDescTS        hlc.Timestamp
+	existingDescName string
+	existingDescID   descpb.ID
+	existingDescTS   hlc.Timestamp
+}
+
+func newMismatchedExternalDataRowTimestampError(
+	newDesc catalog.TableDescriptor, existingDesc catalog.TableDescriptor,
+) *mismatchedExternalDataRowTimestamp {
+	return &mismatchedExternalDataRowTimestamp{
+		newDescName:      newDesc.GetName(),
+		newDescID:        newDesc.GetID(),
+		newDescTS:        newDesc.ExternalRowData().AsOf,
+		existingDescName: existingDesc.GetName(),
+		existingDescID:   existingDesc.GetID(),
+		existingDescTS:   existingDesc.ExternalRowData().AsOf,
+	}
+}
+
+func (e *mismatchedExternalDataRowTimestamp) SafeFormatError(p errors.Printer) (next error) {
+	p.Printf("PCR reader timestamp has moved forward, "+
+		"existing descriptor %s(%d) and timestamp: %s "+
+		"new descritpor %s(%d) and timestamp: %s",
+		e.newDescName,
+		e.newDescID,
+		e.newDescTS,
+		e.existingDescName,
+		e.existingDescID,
+		e.existingDescTS)
+	return nil
+}
+
+func (e *mismatchedExternalDataRowTimestamp) Error() string {
+	return fmt.Sprint(errors.Formattable(e))
+}
+
+var _ errors.SafeFormatter = (*mismatchedExternalDataRowTimestamp)(nil)
+
+// maybeAssertExternalRowDataTS asserts if the descriptor references external
+// row data, then the timestamp across the entire collection *must* match.
+func (ld *leasedDescriptors) maybeAssertExternalRowDataTS(desc catalog.Descriptor) error {
+	tableDesc, ok := desc.(catalog.TableDescriptor)
+	if !ok {
+		return nil
+	}
+	if tableDesc.ExternalRowData() == nil {
+		return nil
+	}
+	currentTS := tableDesc.ExternalRowData().AsOf
+	return ld.cache.IterateByID(func(entry catalog.NameEntry) error {
+		// Skip databases / schemas.
+		if entry.GetParentID() == descpb.InvalidID ||
+			entry.GetParentSchemaID() == descpb.InvalidID ||
+			entry.GetID() == tableDesc.GetID() {
+			return nil
+		}
+		// Next get the underlying descriptor.
+		otherDesc := entry.(lease.LeasedDescriptor).Underlying()
+		if otherTableDesc, ok := otherDesc.(catalog.TableDescriptor); ok {
+			// Skip conventional descriptors.
+			if otherTableDesc.ExternalRowData() == nil {
+				return nil
+			}
+			// Confirm the timestamps match the most recent descriptor.
+			if !otherTableDesc.ExternalRowData().AsOf.Equal(currentTS) {
+				// Normally the PCR catalog reader will run with AOST timestamps,
+				// if during the setup of the connection executor we were able to
+				// lease the system database descriptor and confirm that it is for
+				// a PCR reader catalog. If we were not able to lease the system database
+				// descriptor, then its possible no AOST timestamp is set. Otherwise,
+				// this error should *never* happen.
+				return newMismatchedExternalDataRowTimestampError(
+					tableDesc,
+					otherTableDesc)
+			}
+			// Otherwise, we expect all other timestamps to match as well.
+			return iterutil.StopIteration()
+		}
+		return nil
+	})
+}
+
 // 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,
@@ -99,7 +187,8 @@ func (ld *leasedDescriptors) getByName(
 			log.Eventf(ctx, "found descriptor in collection for (%d, %d, '%s'): %d",
 				parentID, parentSchemaID, name, cached.GetID())
 		}
-		return cached.(lease.LeasedDescriptor).Underlying(), false, nil
+		desc = cached.(lease.LeasedDescriptor).Underlying()
+		return desc, false, nil
 	}
 
 	readTimestamp := txn.ReadTimestamp()
@@ -183,7 +272,12 @@ func (ld *leasedDescriptors) getResult(
 			return nil, false, err
 		}
 	}
-	return ldesc.Underlying(), false, nil
+
+	desc := ldesc.Underlying()
+	if err = ld.maybeAssertExternalRowDataTS(desc); err != nil {
+		return nil, false, err
+	}
+	return desc, false, nil
 }
 
 func (ld *leasedDescriptors) maybeUpdateDeadline(

pkg/sql/catalog/lease/lease.go

@@ -889,6 +889,10 @@ type Manager struct {
 		rangeFeed *rangefeed.RangeFeed
 	}
 
+	// closeTimeStamp for the range feed, which is the timestamp
+	// that we have all the updates for.
+	closeTimestamp atomic.Value
+
 	draining atomic.Value
 
 	// names is a cache for name -> id mappings. A mapping for the cache
@@ -996,6 +1000,12 @@ func NewLeaseManager(
 	lm.storage.writer = newKVWriter(codec, db.KV(), keys.LeaseTableID, settingsWatcher, lm)
 	lm.stopper.AddCloser(lm.sem.Closer("stopper"))
 	lm.mu.descriptors = make(map[descpb.ID]*descriptorState)
+	// We are going to start the range feed later when RefreshLeases
+	// is invoked inside pre-start. So, that guarantees all range feed events
+	// that will be generated will be after the current time. So, historical
+	// queries with in this tenant (i.e. PCR catalog reader) before this point are
+	// guaranteed to be up to date.
+	lm.closeTimestamp.Store(db.KV().Clock().Now())
 	lm.draining.Store(false)
 	lm.descUpdateCh = make(chan catalog.Descriptor)
 	lm.descDelCh = make(chan descpb.ID)
@@ -1438,6 +1448,12 @@ func (m *Manager) RefreshLeases(ctx context.Context, s *stop.Stopper, db *kv.DB)
 	})
 }
 
+// GetSafeReplicationTS gets the timestamp till which the leased descriptors
+// have been synced.
+func (m *Manager) GetSafeReplicationTS() hlc.Timestamp {
+	return m.closeTimestamp.Load().(hlc.Timestamp)
+}
+
 // watchForUpdates will watch a rangefeed on the system.descriptor table for
 // updates.
 func (m *Manager) watchForUpdates(ctx context.Context) {
@@ -1496,6 +1512,7 @@ func (m *Manager) watchForUpdates(ctx context.Context) {
 			return
 		}
 		m.mu.rangeFeedCheckpoints += 1
+		m.closeTimestamp.Store(checkpoint.ResolvedTS)
 	}
 
 	// If we already started a range feed terminate it first

pkg/sql/catalog/replication/BUILD.bazel

@@ -35,10 +35,17 @@ go_test(
         "//pkg/security/securitytest",
         "//pkg/server",
         "//pkg/sql",
+        "//pkg/sql/catalog/lease",
+        "//pkg/testutils",
         "//pkg/testutils/serverutils",
+        "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
+        "//pkg/util/ctxgroup",
+        "//pkg/util/hlc",
+        "//pkg/util/leaktest",
         "//pkg/util/randutil",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )