Merge #61820

61820: sql/catalog/descs: use allDescriptors as a negative cache for by-ID lookups r=ajwerner a=ajwerner

When doing introspection queries we tend to pull and cache the complete
set of descriptors. This powers subsequent calls to various virtual
tables. However, we also often need to grab descriptors by ID. This
happens in some of the virtual indexes and it especially happens in
the call to `pg_table_is_visible`. When that function was called on
all of the oids in pg_class it was forced to make a bunch of kv lookups;
namely one for each and every index and virtual table (depending on
the query plan). We can short-circuit all of that with a map of descriptors
we are already holding in memory.

Before this change, the django introspection query was doing O(indexes) lookups
and now it is doing 2. 

Release note (performance improvement): Reduce the number of round-trips
required to call `pg_table_is_visible` in the context of pg_catalog
queries.

Co-authored-by: Andrew Werner <[email protected]>

pkg/bench/ddl_analysis/testdata/benchmark_expectations

@@ -53,9 +53,9 @@ exp,benchmark
 16,GrantRole/grant_1_role
 19,GrantRole/grant_2_roles
 5,ORMQueries/activerecord_type_introspection_query
-10,ORMQueries/django_table_introspection_1_table
-13,ORMQueries/django_table_introspection_4_tables
-17,ORMQueries/django_table_introspection_8_tables
+2,ORMQueries/django_table_introspection_1_table
+2,ORMQueries/django_table_introspection_4_tables
+2,ORMQueries/django_table_introspection_8_tables
 18,Revoke/revoke_all_on_1_table
 24,Revoke/revoke_all_on_2_tables
 30,Revoke/revoke_all_on_3_tables

pkg/sql/catalog/descs/collection.go

@@ -183,7 +183,7 @@ type Collection struct {
 	//
 	// TODO(ajwerner): This cache may be problematic in clusters with very large
 	// numbers of descriptors.
-	allDescriptors []catalog.Descriptor
+	allDescriptors allDescriptors
 
 	// allDatabaseDescriptors is a slice of all available database descriptors.
 	// These are purged at the same time as allDescriptors.
@@ -223,6 +223,41 @@ type Collection struct {
 	skipValidationOnWrite bool
 }
 
+// allDescriptors is an abstraction to capture the complete set of descriptors
+// read from the store. It is used to accelerate repeated invocations of virtual
+// tables which utilize descriptors. It tends to get used to build a
+// sql.internalLookupCtx.
+//
+// TODO(ajwerner): Memory monitoring.
+// TODO(ajwerner): Unify this struct with the uncommittedDescriptors set.
+// TODO(ajwerner): Unify the sql.internalLookupCtx with the descs.Collection.
+type allDescriptors struct {
+	descs []catalog.Descriptor
+	byID  map[descpb.ID]int
+}
+
+func (d *allDescriptors) init(descriptors []catalog.Descriptor) {
+	d.descs = descriptors
+	d.byID = make(map[descpb.ID]int, len(descriptors))
+	for i, desc := range descriptors {
+		d.byID[desc.GetID()] = i
+	}
+}
+
+func (d *allDescriptors) clear() {
+	d.descs = nil
+	d.byID = nil
+}
+
+func (d *allDescriptors) isEmpty() bool {
+	return d.descs == nil
+}
+
+func (d *allDescriptors) contains(id descpb.ID) bool {
+	_, exists := d.byID[id]
+	return exists
+}
+
 // 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,
@@ -1022,6 +1057,15 @@ func (tc *Collection) getDescriptorByIDMaybeSetTxnDeadline(
 			return readFromStore()
 		}
 
+		// If we have already read all of the descriptor, use it as a negative
+		// cache to short-circuit a lookup we know will be doomed to fail.
+		//
+		// TODO(ajwerner): More generally leverage this set of read descriptors on
+		// the resolution path.
+		if !tc.allDescriptors.isEmpty() && !tc.allDescriptors.contains(id) {
+			return nil, catalog.ErrDescriptorNotFound
+		}
+
 		desc, err := tc.getLeasedDescriptorByID(ctx, txn, id, setTxnDeadline)
 		if err != nil {
 			if errors.Is(err, catalog.ErrDescriptorNotFound) || catalog.HasInactiveDescriptorError(err) {
@@ -1666,7 +1710,7 @@ func (tc *Collection) getUncommittedDescriptorByID(id descpb.ID) *uncommittedDes
 func (tc *Collection) GetAllDescriptors(
 	ctx context.Context, txn *kv.Txn,
 ) ([]catalog.Descriptor, error) {
-	if tc.allDescriptors == nil {
+	if tc.allDescriptors.isEmpty() {
 		descs, err := catalogkv.GetAllDescriptors(ctx, txn, tc.codec())
 		if err != nil {
 			return nil, err
@@ -1680,9 +1724,9 @@ func (tc *Collection) GetAllDescriptors(
 			log.Errorf(ctx, "%s", err.Error())
 		}
 
-		tc.allDescriptors = descs
+		tc.allDescriptors.init(descs)
 	}
-	return tc.allDescriptors, nil
+	return tc.allDescriptors.descs, nil
 }
 
 // HydrateGivenDescriptors installs type metadata in the types present for all
@@ -1902,7 +1946,7 @@ func (tc *Collection) GetObjectNames(
 // releaseAllDescriptors releases the cached slice of all descriptors
 // held by Collection.
 func (tc *Collection) releaseAllDescriptors() {
-	tc.allDescriptors = nil
+	tc.allDescriptors.clear()
 	tc.allDatabaseDescriptors = nil
 	tc.allSchemasForDatabase = nil
 }

pkg/sql/resolver.go

@@ -211,6 +211,8 @@ func (p *planner) CommonLookupFlags(required bool) tree.CommonLookupFlags {
 func (p *planner) IsTableVisible(
 	ctx context.Context, curDB string, searchPath sessiondata.SearchPath, tableID int64,
 ) (isVisible, exists bool, err error) {
+	// TODO(ajwerner): look at this error and only no-op if it is
+	// ErrDescriptorNotFound or something like it.
 	tableDesc, err := p.LookupTableByID(ctx, descpb.ID(tableID))
 	if err != nil {
 		// If an error happened here, it means the table doesn't exist, so we