Merge a11aef2 into blathers/backport-release-25.4-154711

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

@@ -4166,11 +4166,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
+	})
+}