catalog/lease: implement bulk lease prefetch API

Previously, accessing leased descriptors for `crdb_internal` functions
required a round trip per descriptor in cold cases. This resulted in
significant performance overhead.

This change modifies the lease manager to expose an API for prefetching
large numbers of descriptors, enabling bulk acquisition. It updates
descriptor collection to utilize this API in `GetAll.*` operations.
Benchmarks for both cold and warm cases are also included.

Fixes: #158086
Release note: None

Author: fqazi

pkg/sql/catalog/descs/collection.go

@@ -416,6 +416,13 @@ var allowLeasedDescriptorsInCatalogViews = settings.RegisterBoolSetting(
 	settings.WithPublic,
 )
 
+var prefetchLeasedDescriptorsInCatalogViews = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"sql.catalog.allow_leased_descriptors.prefetch.enabled",
+	"if true, catalog views (crdb_internal, information_schema, pg_catalog) can prefetch leased descriptors for improved performance",
+	true,
+)
+
 // GetCatalogGetAllOptions returns the functional options for GetAll* methods
 // based on the cluster setting for allowing leased descriptors in catalog views.
 func GetCatalogGetAllOptions(sv *settings.Values) []GetAllOption {
@@ -1058,6 +1065,11 @@ func (tc *Collection) GetAllObjectsInSchema(
 		if err != nil {
 			return nstree.Catalog{}, err
 		}
+		if options.allowLeased && prefetchLeasedDescriptorsInCatalogViews.Get(&tc.settings.SV) {
+			if err := tc.leased.ensureLeasesExist(ctx, stored.OrderedDescriptorIDs()); err != nil {
+				return nstree.Catalog{}, err
+			}
+		}
 		ret, err = tc.aggregateAllLayers(ctx, txn, options, stored, sc)
 		if err != nil {
 			return nstree.Catalog{}, err
@@ -1084,6 +1096,11 @@ func (tc *Collection) GetAllInDatabase(
 	if err != nil {
 		return nstree.Catalog{}, err
 	}
+	if options.allowLeased && prefetchLeasedDescriptorsInCatalogViews.Get(&tc.settings.SV) {
+		if err := tc.leased.ensureLeasesExist(ctx, stored.OrderedDescriptorIDs()); err != nil {
+			return nstree.Catalog{}, err
+		}
+	}
 	schemas, err := tc.GetAllSchemasInDatabase(ctx, txn, db, opts...)
 	if err != nil {
 		return nstree.Catalog{}, err
@@ -1131,6 +1148,11 @@ func (tc *Collection) GetAllTablesInDatabase(
 	if err != nil {
 		return nstree.Catalog{}, err
 	}
+	if options.allowLeased && prefetchLeasedDescriptorsInCatalogViews.Get(&tc.settings.SV) {
+		if err := tc.leased.ensureLeasesExist(ctx, stored.OrderedDescriptorIDs()); err != nil {
+			return nstree.Catalog{}, err
+		}
+	}
 	var ret nstree.MutableCatalog
 	if db.HasPublicSchemaWithDescriptor() {
 		ret, err = tc.aggregateAllLayers(ctx, txn, options, stored)

pkg/sql/catalog/descs/leased_descriptors.go

@@ -33,6 +33,8 @@ type LeaseManager interface {
 		ctx context.Context, timestamp lease.ReadTimestamp, id descpb.ID,
 	) (lease.LeasedDescriptor, error)
 
+	EnsureBatch(ctx context.Context, ids []descpb.ID) error
+
 	IncGaugeAfterLeaseDuration(
 		gaugeType lease.AfterLeaseDurationGauge,
 	) (decrAfterWait func())
@@ -203,6 +205,13 @@ func (ld *leasedDescriptors) maybeInitReadTimestamp(ctx context.Context, txn dea
 	ld.leaseTimestamp = ld.lm.GetReadTimestamp(ctx, readTimestamp)
 }
 
+// ensureLeasesExist requests that the lease manager acquire leases for the
+// given IDs, if one isn't already exists. This is done in a bulk fashion leading
+// to fewer round trips inside the lease manager.
+func (ld *leasedDescriptors) ensureLeasesExist(ctx context.Context, ids []descpb.ID) error {
+	return ld.lm.EnsureBatch(ctx, ids)
+}
+
 // 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,

pkg/sql/catalog/lease/BUILD.bazel

@@ -101,6 +101,8 @@ go_test(
         "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/roachpb",
+        "//pkg/rpc",
+        "//pkg/rpc/nodedialer",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",

pkg/sql/catalog/lease/descriptor_state.go

@@ -58,6 +58,9 @@ type descriptorState struct {
 		// acquisition finishes but indicate that that new lease is expired are not
 		// ignored.
 		acquisitionsInProgress int
+
+		// acquisitionChannel indicates that a bulk acquisition is in progress.
+		acquisitionChannel chan struct{}
 	}
 }
 

pkg/sql/catalog/lease/lease.go

@@ -108,6 +108,11 @@ var LockedLeaseTimestamp = settings.RegisterBoolSetting(settings.ApplicationLeve
 		"descriptors used can be intentionally older to support this",
 	false)
 
+var MaxBatchLeaseCount = settings.RegisterIntSetting(settings.ApplicationLevel,
+	"sql.catalog.descriptor_lease.max_batch_lease_count",
+	"the maximum number of descriptors to lease in a single batch",
+	1000)
+
 // WaitForNoVersion returns once there are no unexpired leases left
 // for any version of the descriptor.
 func (m *Manager) WaitForNoVersion(
@@ -1009,7 +1014,9 @@ func (m *Manager) insertDescriptorVersions(
 // get needs to see some descriptor updates that we know happened recently.
 func (m *Manager) AcquireFreshestFromStore(ctx context.Context, id descpb.ID) error {
 	// Create descriptorState if needed.
-	_ = m.findDescriptorState(id, true /* create */)
+	state := m.findDescriptorState(id, true /* create */)
+	state.markAcquisitionStart(ctx)
+	defer state.markAcquisitionDone(ctx)
 	// We need to acquire a lease on a "fresh" descriptor, meaning that joining
 	// a potential in-progress lease acquisition is generally not good enough.
 	// If we are to join an in-progress acquisition, it needs to be an acquisition
@@ -1045,13 +1052,159 @@ func (m *Manager) AcquireFreshestFromStore(ctx context.Context, id descpb.ID) er
 	return nil
 }
 
+// upsertDescriptorIntoState inserts a descriptor into the descriptor state.
+func (m *Manager) upsertDescriptorIntoState(
+	ctx context.Context, id descpb.ID, session sqlliveness.Session, desc catalog.Descriptor,
+) error {
+	t := m.findDescriptorState(id, false /* create */)
+	if t == nil {
+		return errors.AssertionFailedf("could not find descriptor state for id %d", id)
+	}
+	t.mu.Lock()
+	t.mu.takenOffline = false
+	defer t.mu.Unlock()
+	err := t.upsertLeaseLocked(ctx, desc, session, m.storage.getRegionPrefix())
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// maybePrepareDescriptorForBulkAcquisition sets up the descriptor state for a bulk
+// acquisition. Returns true if the descriptor has been setup for bulk acquistion,
+// and false if an acqusition is already fetching it.
+func (m *Manager) maybePrepareDescriptorForBulkAcquisition(
+	id descpb.ID, acquisitionCh chan struct{},
+) bool {
+	state := m.findDescriptorState(id, true)
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	// Only work on descriptors that haven't been acquired or have no
+	// acquisition in progress.
+	if (state.mu.takenOffline || len(state.mu.active.data) > 0) ||
+		state.mu.acquisitionsInProgress > 0 {
+		return false
+	}
+	// Mark the descriptor as having an acquisition in progress, any normal
+	// acquisition will wait for the bulk acquisition to finish.
+	state.mu.acquisitionsInProgress++
+	state.mu.acquisitionChannel = acquisitionCh
+	return true
+}
+
+// completeBulkAcquisition clears up the descriptor state after a bulk acquisition.
+func (m *Manager) completeBulkAcquisition(id descpb.ID) {
+	state := m.findDescriptorState(id, false)
+	if state == nil {
+		return
+	}
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	state.mu.acquisitionChannel = nil
+	state.mu.acquisitionsInProgress -= 1
+}
+
+// EnsureBatch ensures that a set of IDs are already cached by the lease manager
+// by doing a bulk acquisition. The acquisition will be done on a best effort
+// basis, where missing descriptors, dropped, adding, or validation errors will
+// be ignored.
+func (m *Manager) EnsureBatch(ctx context.Context, ids []descpb.ID) error {
+	maxBatchSize := MaxBatchLeaseCount.Get(&m.storage.settings.SV)
+	idsToFetch := make([]descpb.ID, 0, min(len(ids), int(maxBatchSize)))
+	lastVersion := make([]descpb.DescriptorVersion, cap(idsToFetch))
+	lastSession := make([]sqlliveness.SessionID, cap(idsToFetch))
+	for len(ids) > 0 {
+		processNextBatch := func() error {
+			// Reset for the next batch.
+			idsToFetch = idsToFetch[:0]
+			acquisitionCh := make(chan struct{})
+			defer close(acquisitionCh)
+
+			// Figure out which IDs have no state object allocated.
+			batchCompleted := true
+			for idx, id := range ids {
+				if !m.maybePrepareDescriptorForBulkAcquisition(id, acquisitionCh) {
+					continue
+				}
+				idsToFetch = append(idsToFetch, id)
+				if len(idsToFetch) == int(maxBatchSize) {
+					ids = ids[idx+1:]
+					batchCompleted = false
+					break
+				}
+			}
+			// All of the entries in the slice have been processed.
+			if batchCompleted {
+				ids = nil
+			}
+			// Nothing needs to be fetched everything is cached in memory.
+			if len(idsToFetch) == 0 {
+				return nil
+			}
+			// Clean up the acquisition state after-wards.
+			defer func() {
+				for _, id := range idsToFetch {
+					m.completeBulkAcquisition(id)
+				}
+			}()
+			// Initial acquisition so version and session ID are blank for any
+			// previous version.
+			lastVersion = lastVersion[:len(idsToFetch)]
+			lastSession = lastSession[:len(idsToFetch)]
+			// Execute a batch acquisition at the storage layer.
+			session, err := m.storage.livenessProvider.Session(ctx)
+			if err != nil {
+				return errors.Wrapf(err, "lease acquisition was unable to resolve liveness session")
+			}
+			batch, _, err := m.storage.acquireBatch(ctx, session, idsToFetch, lastVersion, lastSession)
+			if err != nil {
+				return err
+			}
+			// Upsert the descriptors into the descriptor state.
+			for idx, id := range idsToFetch {
+				// If any descriptors in the batch have failed to acquire, we are going
+				// to skip them. This can be due to validation error, being dropped, or
+				// being added, or missing. The caller can surface these when it attempts
+				// to access these leases, and we will attempt acquisition again.
+				if batch.errs[idx] != nil {
+					continue
+				}
+				if err := m.upsertDescriptorIntoState(ctx, id, session, batch.descs[idx]); err != nil {
+					return err
+				}
+			}
+			return nil
+		}
+		if err := processNextBatch(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // If the lease cannot be obtained because the descriptor is in the process of
 // being dropped or offline, the error will be of type inactiveTableError.
 // The boolean returned is true if this call was actually responsible for the
-// lease acquisition.
+// lease acquisition. Any callers must allocate the descriptor state and mark
+// that an acquisition is in progress with descriptorState.markAcquisitionStart,
+// before invoking this method.
 func (m *Manager) acquireNodeLease(
 	ctx context.Context, id descpb.ID, typ AcquireType,
 ) (bool, error) {
+	// Precondition: Ensure that the descriptor state is setup and ready for any
+	// acquisition.
+	if buildutil.CrdbTestBuild {
+		state := m.findDescriptorState(id, false)
+		if state == nil {
+			return false, errors.AssertionFailedf("descriptor state must be allocated before acquistion of %d", id)
+		}
+		state.mu.Lock()
+		acquisition := state.mu.acquisitionsInProgress
+		state.mu.Unlock()
+		if acquisition <= 0 {
+			return false, errors.AssertionFailedf("descriptor acquistion must be marked before invoking acquire node lease on %d.", id)
+		}
+	}
 	start := timeutil.Now()
 	log.VEventf(ctx, 2, "acquiring lease for descriptor %d...", id)
 	future, didAcquire := m.storage.group.DoChan(ctx,
@@ -1064,6 +1217,25 @@ func (m *Manager) acquireNodeLease(
 			if m.IsDraining() {
 				return false, errLeaseManagerIsDraining
 			}
+			waitedForBulkAcquire, err := func() (bool, error) {
+				state := m.findDescriptorState(id, false)
+				state.mu.Lock()
+				acquisitionChannel := state.mu.acquisitionChannel
+				state.mu.Unlock()
+				if acquisitionChannel == nil {
+					return false, nil
+				}
+				select {
+				case <-acquisitionChannel:
+					return true, nil
+				case <-ctx.Done():
+					return false, ctx.Err()
+				}
+			}()
+			if waitedForBulkAcquire || err != nil {
+				// We didn't do any acquiring and just waited for a bulk operation.
+				return false, err
+			}
 			newest := m.findNewest(id)
 			var currentVersion descpb.DescriptorVersion
 			var currentSessionID sqlliveness.SessionID
@@ -1086,22 +1258,6 @@ func (m *Manager) acquireNodeLease(
 				return desc, nil
 			}
 
-			doUpsertion := func(desc catalog.Descriptor) error {
-				t := m.findDescriptorState(id, false /* create */)
-				if t == nil {
-					return errors.AssertionFailedf("could not find descriptor state for id %d", id)
-				}
-				t.mu.Lock()
-				t.mu.takenOffline = false
-				defer t.mu.Unlock()
-				err = t.upsertLeaseLocked(ctx, desc, session, m.storage.getRegionPrefix())
-				if err != nil {
-					return err
-				}
-
-				return nil
-			}
-
 			// These tables are special and can have their versions bumped
 			// without blocking on other nodes converging to that version.
 			if newest != nil && (id == keys.UsersTableID ||
@@ -1117,8 +1273,7 @@ func (m *Manager) acquireNodeLease(
 				if err != nil {
 					return false, err
 				}
-				err = doUpsertion(desc)
-				if err != nil {
+				if err := m.upsertDescriptorIntoState(ctx, id, session, desc); err != nil {
 					return false, err
 				}
 
@@ -1139,8 +1294,7 @@ func (m *Manager) acquireNodeLease(
 				if desc == nil {
 					return true, nil
 				}
-				err = doUpsertion(desc)
-				if err != nil {
+				if err := m.upsertDescriptorIntoState(ctx, id, session, desc); err != nil {
 					return false, err
 				}
 			}
@@ -2952,6 +3106,10 @@ func (m *Manager) refreshSomeLeases(ctx context.Context, refreshAndPurgeAllDescr
 						return
 					}
 				}
+				// Mark that an acquisition is in progress.
+				state := m.findDescriptorState(id, false)
+				state.markAcquisitionStart(ctx)
+				defer state.markAcquisitionDone(ctx)
 				if _, err := m.acquireNodeLease(ctx, id, AcquireBackground); err != nil {
 					log.Dev.Errorf(ctx, "refreshing descriptor: %d lease failed: %s", id, err)
 

pkg/sql/catalog/lease/lease_internal_test.go

@@ -573,12 +573,15 @@ CREATE TABLE t.%s (k CHAR PRIMARY KEY, v CHAR);
 	}
 	expiration := lease.Expiration(context.Background())
 	// Acquire another lease.
-	if _, err := leaseManager.acquireNodeLease(
-		context.Background(), tableDesc.GetID(), AcquireBlock,
-	); err != nil {
-		t.Fatal(err)
-	}
-
+	func() {
+		lease.t.markAcquisitionStart(context.Background())
+		defer lease.t.markAcquisitionDone(context.Background())
+		if _, err := leaseManager.acquireNodeLease(
+			context.Background(), tableDesc.GetID(), AcquireBlock,
+		); err != nil {
+			t.Fatal(err)
+		}
+	}()
 	// Check the name resolves to the new lease.
 	newLease, _ := leaseManager.names.get(context.Background(), tableDesc.GetParentID(), tableDesc.GetParentSchemaID(), tableName, s.Clock().Now())
 	if newLease == nil {