Merge #158087

158087: catalog/lease: add support for bulk leasing r=fqazi a=fqazi

This patch does the following to help support bulk leasing in the crdb_internal / pg_catalog / information_schema views:

1. Refactor the storage layer for the lease manager to support bulk acquisitions by introducing a new `acquireBatch` API, which allows the lease manager to acquire multiple leases in a transaction. This also is used to implement `acquire`
2. Add a new `EnsureBatch` public method in the lease manager, which allows the lease manager to acquire multiple leases into memory, which acts as a prefetch mechanism. If a normal acquisition occurs concurrently it will correctly wait for the bulk acquisition to complete.
3. Modification to descriptor collection API to use EnsureBatch in the GetAll.* API.


New benchmarks for these change and leased descriptors vs KV in GetAll.*:

```
Cold
======================
BenchmarkLargeDatabaseColdPgClass/pg_class/leased_descriptors=true/use_prefetch=true/tables=8192         	1000000000	         0.7835 ns/op	       0 B/op	       0 allocs/op
BenchmarkLargeDatabaseColdPgClass/pg_class/leased_descriptors=true/use_prefetch=false/tables=8192        	       1	8642132499 ns/op	1731154008 B/op	12307123 allocs/op
BenchmarkLargeDatabaseColdPgClass/pg_class/leased_descriptors=false/use_prefetch=false/tables=8192       	1000000000	         0.3063 ns/op	       0 B/op	       0 allocs/op

Warm
======================
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=true/tables=8192         	       6	 170890458 ns/op	90846753 B/op	  428243 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=false/tables=8192        	       7	 192905720 ns/op	91194598 B/op	  431132 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=false/use_prefetch=false/tables=8192       	       4	 269645802 ns/op	243010064 B/op	 1401622 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=true/tables=16384        	       3	 352242778 ns/op	182838792 B/op	  847251 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=false/tables=16384       	       3	 399849806 ns/op	182671565 B/op	  847271 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=false/use_prefetch=false/tables=16384      	       2	 582772520 ns/op	488168660 B/op	 2793079 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=true/tables=32768        	       2	 778396542 ns/op	365770676 B/op	 1690881 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=true/use_prefetch=false/tables=32768       	       2	 803655520 ns/op	364708036 B/op	 1685706 allocs/op
BenchmarkLargeDatabaseWarmPgClass/pg_class/leased_descriptors=false/use_prefetch=false/tables=32768      	       1	1300587666 ns/op	977227432 B/op	 5585049 allocs/tp

```

Co-authored-by: Faizan Qazi <[email protected]>

Author: craig[bot]

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/ie_writer_test.go

@@ -61,4 +61,13 @@ func (w *ieWriter) insertLease(ctx context.Context, txn *kv.Txn, l leaseFields)
 	return nil
 }
 
+func (w *ieWriter) insertLeases(ctx context.Context, txn *kv.Txn, leases []leaseFields) error {
+	for _, l := range leases {
+		if err := w.insertLease(ctx, txn, l); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 var _ writer = (*ieWriter)(nil)

pkg/sql/catalog/lease/kv_writer.go

@@ -52,23 +52,29 @@ func leaseTableWithID(id descpb.ID, table systemschema.SystemTable) catalog.Tabl
 	return mut.ImmutableCopy().(catalog.TableDescriptor)
 }
 
-func (w *kvWriter) insertLease(ctx context.Context, txn *kv.Txn, l leaseFields) error {
-	if err := w.do(ctx, txn, l, func(b *kv.Batch) error {
-		if l.sessionID != nil {
-			err := w.sessionBasedWriter.Insert(ctx, b, false /*kvTrace*/, leaseAsSessionBasedDatum(l)...)
-			if err != nil {
-				return err
+func (w *kvWriter) insertLeases(ctx context.Context, txn *kv.Txn, leases []leaseFields) error {
+	if err := w.do(ctx, txn, leases, func(b *kv.Batch) error {
+		for _, l := range leases {
+			if l.sessionID != nil {
+				err := w.sessionBasedWriter.Insert(ctx, b, false /*kvTrace*/, leaseAsSessionBasedDatum(l)...)
+				if err != nil {
+					return err
+				}
 			}
 		}
 		return nil
 	}); err != nil {
-		return errors.Wrapf(err, "failed to insert lease %v", l)
+		return errors.Wrapf(err, "failed to insert lease %v", leases)
 	}
 	return nil
 }
 
+func (w *kvWriter) insertLease(ctx context.Context, txn *kv.Txn, l leaseFields) error {
+	return w.insertLeases(ctx, txn, []leaseFields{l})
+}
+
 func (w *kvWriter) deleteLease(ctx context.Context, txn *kv.Txn, l leaseFields) error {
-	if err := w.do(ctx, txn, l, func(b *kv.Batch) error {
+	if err := w.do(ctx, txn, []leaseFields{l}, func(b *kv.Batch) error {
 		if l.sessionID != nil {
 			err := w.sessionBasedWriter.Delete(ctx, b, false /*kvTrace*/, leaseAsSessionBasedDatum(l)...)
 			if err != nil {
@@ -85,7 +91,7 @@ func (w *kvWriter) deleteLease(ctx context.Context, txn *kv.Txn, l leaseFields)
 type addToBatchFunc = func(*kv.Batch) error
 
 func (w *kvWriter) do(
-	ctx context.Context, txn *kv.Txn, lease leaseFields, addToBatch addToBatchFunc,
+	ctx context.Context, txn *kv.Txn, lease []leaseFields, addToBatch addToBatchFunc,
 ) error {
 	run := (*kv.Txn).Run
 	do := func(ctx context.Context, txn *kv.Txn) error {

pkg/sql/catalog/lease/kv_writer_test.go

@@ -282,4 +282,11 @@ func (t teeWriter) insertLease(ctx context.Context, txn *kv.Txn, fields leaseFie
 	)
 }
 
+func (t teeWriter) insertLeases(ctx context.Context, txn *kv.Txn, leases []leaseFields) error {
+	return errors.CombineErrors(
+		t.a.insertLeases(ctx, txn, leases),
+		t.b.insertLeases(ctx, txn, leases),
+	)
+}
+
 var _ writer = (*teeWriter)(nil)