Merge #119140

119140: pkg/sql: export sql.aggregated_livebytes metric for tenants r=jaylim-crl a=jaylim-crl

#### pkg/util/metric: support metrics removal from the metrics registry 

Previously, once a metric has been added to the metrics registry, it will
always be registered forever, and there isn't a mechanism to remove it. For
multi-tenancy, we plan to implement a job that exports global metrics for
tenants (i.e. such metrics should only exist on one SQL node at any point in
time). Given that jobs can be cancelled and resumed on a different SQL node,
the only option to support such a behavior is to remove metrics from the
registry when the job is no longer running, and this commit adds such support
to it.

Epic: none

Release note: None

#### pkg/server: add ApproximateTotalStats field to SpanStats proto message 

Previously, the SpanStats proto message only kept track of the logical MVCC
stats in the TotalStats field. This is insufficient for the work that exposes
the aggregated livebytes as a metric for tenants as the metric value needs to
take into account all replicas for a given range. To address that, this commit
adds a new ApproximateTotalStats field to the SpanStats proto message, and it
represents post-replicated MVCC stats for the span.

Epic: none

Release note: None

#### pkg/sql: export sql.aggregated_livebytes metric for out-of-process tenants

Previously, in order to obtain livebytes metrics for tenants, one would need
to query such values via the KV servers, and this can be problematic if we
only have access to just the SQL servers. For example, in CockroachDB Cloud,
only metrics from the SQL servers are exported to end-users, and is done so
directly from the cockroachdb process. It is not trivial to export an
additional subset of metrics from the KV servers filtered by tenant ID.

To address that, this commit exposes livebytes for tenants directly via an
aggregated metric on the SQL nodes. The aggregated metric will be updated
every 60 seconds by default, and will be exported via the existing MVCC
statistics update job. Unlike other job metrics where metrics are registered
at initialization time and stays forever, this aggregated metric is tied to
the lifespan of the job (i.e. it is only exported if the job is running, and
unexported otherwise).

This feature is scoped to standalone SQL servers only, which at this point of
writing, is only supported in CockroachDB Cloud. If we wanted to backport this
into 23.2, it should be straightforward as well since the permanent upgrade
to insert the job is already in release-23.2.

Fixes: #119139

Epic: none

Release note (sql change): Out-of-process SQL servers will start exporting a
new sql.aggregated_livebytes metric. This metric gets updated once every 60
seconds by default, and its update interval can be configured via the
`tenant_global_metrics_exporter_interval` cluster setting.

Co-authored-by: Jay <[email protected]>

Author: craig[bot]

pkg/jobs/registry.go

@@ -321,6 +321,8 @@ const (
 	// SqlActivityUpdaterJobID A static job ID is used for the SQL activity tables.
 	SqlActivityUpdaterJobID = jobspb.JobID(103)
 
+	// MVCCStatisticsJobID A static job ID used for the MVCC statistics update
+	// job.
 	MVCCStatisticsJobID = jobspb.JobID(104)
 )
 

pkg/roachpb/span_stats.proto

@@ -40,6 +40,7 @@ message SpanStatsRequest {
 message SpanStats {
   // TotalStats are the logical MVCC stats for the requested span.
   cockroach.storage.enginepb.MVCCStats total_stats = 1 [(gogoproto.nullable) = false];
+
   // range_count measures the number of ranges that the request span falls within.
   // A SpanStatsResponse for a span that lies within a range, and whose start
   // key sorts after the range start, and whose end key sorts before the
@@ -63,7 +64,13 @@ message SpanStats {
   // It represents a physical value across all replicas.
   uint64 external_file_bytes = 6;
 
-  // NEXT ID: 7.
+  // ApproximateTotalStats is the approximate, post-replicated MVCC stats for
+  // the requested span. This is an approximate value as it is constructed by
+  // querying all replicas, which includes followers, and the value can be
+  // stale.
+  cockroach.storage.enginepb.MVCCStats approximate_total_stats = 7 [(gogoproto.nullable) = false];
+
+  // NEXT ID: 8.
 }
 
 message SpanStatsResponse {

pkg/server/server_sql.go

@@ -2086,3 +2086,8 @@ func (s *SQLServer) ExecutorConfig() *sql.ExecutorConfig {
 func (s *SQLServer) InternalExecutor() isql.Executor {
 	return s.internalExecutor
 }
+
+// MetricsRegistry returns the application-level metrics registry.
+func (s *SQLServer) MetricsRegistry() *metric.Registry {
+	return s.metricsRegistry
+}

pkg/server/span_stats_server.go

@@ -138,13 +138,16 @@ func (s *systemStatusServer) spanStatsFanOut(
 
 		nodeResponse := resp.(*roachpb.SpanStatsResponse)
 
-		// Values of ApproximateDiskBytes, RemoteFileBytes, and ExternalFileBytes should be physical values, but
-		// TotalStats (MVCC stats) should be the logical, pre-replicated value.
+		// Values of ApproximateTotalStats, ApproximateDiskBytes,
+		// RemoteFileBytes, and ExternalFileBytes should be physical values, but
+		// TotalStats should be the logical, pre-replicated value.
+		//
 		// Note: This implementation can return arbitrarily stale values, because instead of getting
 		// MVCC stats from the leaseholder, MVCC stats are taken from the node that responded first.
 		// See #108779.
 		for spanStr, spanStats := range nodeResponse.SpanToStats {
 			// Accumulate physical values across all replicas:
+			res.SpanToStats[spanStr].ApproximateTotalStats.Add(spanStats.TotalStats)
 			res.SpanToStats[spanStr].ApproximateDiskBytes += spanStats.ApproximateDiskBytes
 			res.SpanToStats[spanStr].RemoteFileBytes += spanStats.RemoteFileBytes
 			res.SpanToStats[spanStr].ExternalFileBytes += spanStats.ExternalFileBytes

pkg/server/span_stats_test.go

@@ -186,6 +186,17 @@ func TestSpanStatsFanOut(t *testing.T) {
 					spanStats.TotalStats.LiveCount,
 				)
 			}
+
+			approxKeys := numNodes * tcase.expectedKeys
+			if approxKeys != spanStats.ApproximateTotalStats.LiveCount {
+				return errors.Newf(
+					"Multi-span: expected %d post-replicated keys in span [%s - %s], found %d",
+					approxKeys,
+					rSpan.Key.String(),
+					rSpan.EndKey.String(),
+					spanStats.TotalStats.LiveCount,
+				)
+			}
 		}
 
 		return nil

pkg/server/status/recorder.go

@@ -223,6 +223,13 @@ func (mr *MetricsRecorder) RemoveTenantRegistry(tenantID roachpb.TenantID) {
 	delete(mr.mu.tenantRegistries, tenantID)
 }
 
+// AppRegistry returns the metric registry for application-level metrics.
+func (mr *MetricsRecorder) AppRegistry() *metric.Registry {
+	mr.mu.Lock()
+	defer mr.mu.Unlock()
+	return mr.mu.appRegistry
+}
+
 // AddNode adds various metric registries an initialized server, along
 // with its descriptor and start time.
 // The registries are:

pkg/sql/BUILD.bazel

@@ -673,6 +673,7 @@ go_test(
         "multitenant_admin_function_test.go",
         "mutation_test.go",
         "mvcc_backfiller_test.go",
+        "mvcc_statistics_update_job_test.go",
         "normalization_test.go",
         "pg_metadata_test.go",
         "pg_oid_test.go",
@@ -787,6 +788,7 @@ go_test(
         "//pkg/server/serverpb",
         "//pkg/server/settingswatcher",
         "//pkg/server/srvtestutils",
+        "//pkg/server/status",
         "//pkg/server/status/statuspb",
         "//pkg/server/telemetry",
         "//pkg/settings",
@@ -883,6 +885,7 @@ go_test(
         "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
+        "//pkg/ts",
         "//pkg/upgrade/upgradebase",
         "//pkg/util",
         "//pkg/util/admission",
@@ -935,6 +938,7 @@ go_test(
         "@com_github_lib_pq//oid",
         "@com_github_petermattis_goid//:goid",
         "@com_github_pmezard_go_difflib//difflib",
+        "@com_github_prometheus_common//expfmt",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@in_gopkg_yaml_v2//:yaml_v2",

pkg/sql/exec_util.go

@@ -1158,10 +1158,11 @@ type NodeInfo struct {
 	PGURL func(*url.Userinfo) (*pgurl.URL, error)
 }
 
-// nodeStatusGenerator is a limited portion of the status.MetricsRecorder
+// limitedMetricsRecorder is a limited portion of the status.MetricsRecorder
 // struct, to avoid having to import all of status in sql.
-type nodeStatusGenerator interface {
+type limitedMetricsRecorder interface {
 	GenerateNodeStatus(ctx context.Context) *statuspb.NodeStatus
+	AppRegistry() *metric.Registry
 }
 
 // SystemTenantOnly wraps an object in the ExecutorConfig that is only
@@ -1230,7 +1231,7 @@ type ExecutorConfig struct {
 	// available when not running as a system tenant.
 	SQLStatusServer    serverpb.SQLStatusServer
 	TenantStatusServer serverpb.TenantStatusServer
-	MetricsRecorder    nodeStatusGenerator
+	MetricsRecorder    limitedMetricsRecorder
 	SessionRegistry    *SessionRegistry
 	ClosedSessionCache *ClosedSessionCache
 	SQLLiveness        sqlliveness.Provider

pkg/sql/mvcc_statistics_update_job.go

@@ -17,40 +17,141 @@ package sql
 
 import (
 	"context"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/metric"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 )
 
+// TenantGlobalMetricsExporterInterval is the interval at which an external
+// tenant's process in the cluster will update the global metrics. This is
+// exported for testing purposes.
+var TenantGlobalMetricsExporterInterval = settings.RegisterDurationSetting(
+	settings.ApplicationLevel,
+	"tenant_global_metrics_exporter_interval",
+	"the interval at which a node in the cluster will update the exported global metrics",
+	60*time.Second,
+	settings.PositiveDuration,
+)
+
+// mvccStatisticsUpdateJob is a singleton job that is meant to update MVCC
+// statistics. Historically, this was added to update system.mvcc_statistics,
+// but the project was deprioritized. Currently, this is used by external
+// process tenants to export global metrics periodically. For such metrics,
+// they will only be present on a SQL node if the job is running. Once the job
+// stops, the metrics will be removed from the metric registry.
 type mvccStatisticsUpdateJob struct {
 	job *jobs.Job
+	st  *cluster.Settings
+
+	// dynamicMetrics keep track of metrics which are added/removed dynamically
+	// as the job runs. Unlike regular job metrics (i.e. WithJobMetrics), which
+	// are registered when the job starts the first time, and never removed from
+	// the metric registry, metrics in this list should be removed when the job
+	// is not running.
+	dynamicMetrics struct {
+		livebytes *metric.Gauge
+	}
 }
 
 var _ jobs.Resumer = (*mvccStatisticsUpdateJob)(nil)
 
-func (j *mvccStatisticsUpdateJob) Resume(ctx context.Context, execCtxI interface{}) (jobErr error) {
+// Resume implements the jobs.Resumer interface.
+func (j *mvccStatisticsUpdateJob) Resume(ctx context.Context, execCtxI interface{}) error {
+	log.Infof(ctx, "starting mvcc statistics update job")
+
+	// This job is a forever running background job, and it is always safe to
+	// terminate the SQL pod whenever the job is running, so mark it as idle.
+	j.job.MarkIdle(true)
+
+	execCtx := execCtxI.(JobExecContext)
+
+	// Export global metrics for tenants if this is an out-of-process SQL node.
+	// All external mode tenant servers have no node IDs.
+	if _, hasNodeID := execCtx.ExecCfg().NodeInfo.NodeID.OptionalNodeID(); !hasNodeID {
+		return j.runTenantGlobalMetricsExporter(ctx, execCtx)
+	}
+
 	// TODO(zachlite):
 	// Delete samples older than configurable setting...
 	// Collect span stats for tenant descriptors...
 	// Write new samples...
-	execCtx := execCtxI.(JobExecContext)
-	stopper := execCtx.ExecCfg().DistSQLSrv.Stopper
-	j.job.MarkIdle(true)
+
+	// Block until context is cancelled since there's nothing that needs to be
+	// done here. We should not return nil, or else the job will be marked as
+	// succeeded.
+	<-ctx.Done()
+	return ctx.Err()
+}
+
+// runTenantGlobalMetricsExporter executes the logic to export global metrics
+// for tenants.
+func (j *mvccStatisticsUpdateJob) runTenantGlobalMetricsExporter(
+	ctx context.Context, execCtx JobExecContext,
+) error {
+	metricsRegistry := execCtx.ExecCfg().MetricsRecorder.AppRegistry()
+
+	initialRun := true
+	defer func() {
+		metricsRegistry.RemoveMetric(j.dynamicMetrics.livebytes)
+	}()
+
+	runTask := func() error {
+		resp, err := execCtx.ExecCfg().TenantStatusServer.SpanStats(
+			ctx,
+			&roachpb.SpanStatsRequest{
+				// Fan out to all nodes. SpanStats takes care of only contacting
+				// the relevant nodes with the tenant's span.
+				NodeID: "0",
+				Spans:  []roachpb.Span{execCtx.ExecCfg().Codec.TenantSpan()},
+			},
+		)
+		if err != nil {
+			return err
+		}
+		var total int64
+		for _, stats := range resp.SpanToStats {
+			total += stats.ApproximateTotalStats.LiveBytes
+		}
+		j.dynamicMetrics.livebytes.Update(total)
+
+		// Only register metrics once we get our initial values. This avoids
+		// metrics from fluctuating whenever the job restarts.
+		if initialRun {
+			metricsRegistry.AddMetric(j.dynamicMetrics.livebytes)
+			initialRun = false
+		}
+		return nil
+	}
+
+	timer := timeutil.NewTimer()
+	defer timer.Stop()
+
+	// Fire the timer immediately to start the initial update.
+	timer.Reset(0)
 
 	for {
 		select {
 		case <-ctx.Done():
-			return nil
-
-		case <-stopper.ShouldQuiesce():
-			return nil
+			return ctx.Err()
+		case <-timer.C:
+			timer.Read = true
+			timer.Reset(TenantGlobalMetricsExporterInterval.Get(&execCtx.ExecCfg().Settings.SV))
+			if err := runTask(); err != nil {
+				log.Errorf(ctx, "mvcc statistics update job error: %v", err)
+			}
 		}
 	}
 }
 
+// OnFailOrCancel implements the jobs.Resumer interface.
 func (j *mvccStatisticsUpdateJob) OnFailOrCancel(
 	ctx context.Context, _ interface{}, jobErr error,
 ) error {
@@ -63,14 +164,23 @@ func (j *mvccStatisticsUpdateJob) OnFailOrCancel(
 	return nil
 }
 
+// CollectProfile implements the jobs.Resumer interface.
 func (j *mvccStatisticsUpdateJob) CollectProfile(_ context.Context, _ interface{}) error {
 	return nil
 }
 
 func init() {
-	jobs.RegisterConstructor(jobspb.TypeMVCCStatisticsUpdate,
+	jobs.RegisterConstructor(
+		jobspb.TypeMVCCStatisticsUpdate,
 		func(job *jobs.Job, settings *cluster.Settings) jobs.Resumer {
-			return &mvccStatisticsUpdateJob{job: job}
+			exporter := &mvccStatisticsUpdateJob{job: job, st: settings}
+			exporter.dynamicMetrics.livebytes = metric.NewGauge(metric.Metadata{
+				Name:        "sql.aggregated_livebytes",
+				Help:        "Aggregated number of bytes of live data (keys plus values)",
+				Measurement: "Storage",
+				Unit:        metric.Unit_BYTES,
+			})
+			return exporter
 		},
 		jobs.DisablesTenantCostControl,
 	)