schematelemetry: update object count gauge using table stats

In order to support larger scales of object counts, we switch away from
using a full table scan on system.descriptors in order to update the
object count gauge.

Instead, we use table stats on system.descriptor now. The
schematelemetry job is updated so it notifies the stats refresher to
keep the stats up to date. Note that the stats refresher is also
notified with a partial count already.

Release note: None

Author: rafiss

pkg/sql/catalog/schematelemetry/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     deps = [
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
+        "//pkg/keys",
         "//pkg/scheduledjobs",
         "//pkg/security/username",
         "//pkg/server/telemetry",

pkg/sql/catalog/schematelemetry/schema_telemetry_job.go

@@ -7,9 +7,11 @@ package schematelemetry
 
 import (
 	"context"
+	"math"
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql"
@@ -63,6 +65,31 @@ func (t schemaTelemetryResumer) Resume(ctx context.Context, execCtx interface{})
 	if k := p.ExecCfg().SchemaTelemetryTestingKnobs; k != nil {
 		knobs = *k
 	}
+	// Notify the stats refresher to update the system.descriptors table stats,
+	// and update the object count in schema changer metrics.
+	err := p.ExecCfg().InternalDB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
+		desc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.DescriptorTableID)
+		if err != nil {
+			return err
+		}
+		p.ExecCfg().StatsRefresher.NotifyMutation(desc, math.MaxInt64 /* rowCount */)
+
+		// Note: This won't be perfectly up-to-date, but it will make sure the
+		// metric gets updated periodically. It also gets updated after every
+		// schema change.
+		tableStats, err := p.ExecCfg().TableStatsCache.GetTableStats(ctx, desc, nil /* typeResolver */)
+		if err != nil {
+			return err
+		}
+		if len(tableStats) > 0 {
+			// Use the row count from the most recent statistic.
+			p.ExecCfg().SchemaChangerMetrics.ObjectCount.Update(int64(tableStats[0].RowCount))
+		}
+		return nil
+	})
+	if err != nil {
+		return errors.Wrap(err, "failed to notify stats refresher to update system.descriptors table stats")
+	}
 
 	// Outside of tests, scan the catalog tables AS OF SYSTEM TIME slightly in the
 	// past. Schema telemetry is not latency-sensitive to the point where a few

pkg/sql/conn_executor.go

@@ -4175,11 +4175,10 @@ func (ex *connExecutor) txnStateTransitionsApplyWrapper(
 			if err := ex.waitForInitialVersionForNewDescriptors(cachedRegions); err != nil {
 				return advanceInfo{}, err
 			}
-		}
-		if ex.extraTxnState.descCollection.CountUncommittedNewOrDroppedDescriptors() > 0 {
+
 			execCfg := ex.planner.ExecCfg()
 			if err := UpdateDescriptorCount(ex.Ctx(), execCfg, execCfg.SchemaChangerMetrics); err != nil {
-				log.Dev.Warningf(ex.Ctx(), "failed to scan descriptor table: %v", err)
+				log.Dev.Warningf(ex.Ctx(), "failed to update descriptor count metric: %v", err)
 			}
 		}
 		fallthrough

pkg/sql/schema_changer.go

@@ -3574,20 +3574,3 @@ func (p *planner) CanCreateCrossDBSequenceRef() error {
 	}
 	return nil
 }
-
-// UpdateDescriptorCount updates our sql.schema_changer.object_count gauge with
-// a fresh count of objects in the system.descriptor table.
-func UpdateDescriptorCount(
-	ctx context.Context, execCfg *ExecutorConfig, metric *SchemaChangerMetrics,
-) error {
-	return DescsTxn(ctx, execCfg, func(ctx context.Context, txn isql.Txn, col *descs.Collection) error {
-		row, err := txn.QueryRow(ctx, "sql-schema-changer-object-count", txn.KV(),
-			`SELECT count(*) FROM system.descriptor`)
-		if err != nil {
-			return err
-		}
-		count := *row[0].(*tree.DInt)
-		metric.ObjectCount.Update(int64(count))
-		return nil
-	})
-}

pkg/sql/schema_changer_metrics.go

@@ -6,7 +6,11 @@
 package sql
 
 import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 )
@@ -42,3 +46,25 @@ func NewSchemaChangerMetrics() *SchemaChangerMetrics {
 		ObjectCount:         metric.NewGauge(metaObjects),
 	}
 }
+
+// UpdateDescriptorCount updates our sql.schema_changer.object_count gauge with
+// a fresh count of objects in the system.descriptor table.
+func UpdateDescriptorCount(
+	ctx context.Context, execCfg *ExecutorConfig, metric *SchemaChangerMetrics,
+) error {
+	return execCfg.InternalDB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
+		desc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.DescriptorTableID)
+		if err != nil {
+			return err
+		}
+		tableStats, err := execCfg.TableStatsCache.GetTableStats(ctx, desc, nil /* typeResolver */)
+		if err != nil {
+			return err
+		}
+		if len(tableStats) > 0 {
+			// Use the row count from the most recent statistic.
+			metric.ObjectCount.Update(int64(tableStats[0].RowCount))
+		}
+		return nil
+	})
+}