feat: metrics module

cmd/spicedb/memoryprotection/memory_protection_integration_test.go

@@ -33,9 +33,18 @@ func TestServeWithMemoryProtectionMiddleware(t *testing.T) {
 
 	serverToken := "mykey"
 	serveResource, err := pool.RunWithOptions(&dockertest.RunOptions{
-		Repository:   "authzed/spicedb",
-		Tag:          "ci",
-		Cmd:          []string{"serve", "--log-level=debug", "--grpc-preshared-key", serverToken, "--telemetry-endpoint=\"\""},
+		Repository: "authzed/spicedb",
+		Tag:        "ci",
+		Cmd: []string{
+			"serve",
+			"--log-level=debug",
+			"--grpc-preshared-key", serverToken,
+			"--telemetry-endpoint", "",
+			// With very low GOMEMLIMIT values, percentage-based dispatch cache defaults
+			// can round down to zero and fail startup.
+			"--dispatch-cache-max-cost", "8KiB",
+			"--dispatch-cluster-cache-max-cost", "8KiB",
+		},
 		ExposedPorts: []string{"50051/tcp"},
 		Env:          []string{"GOMEMLIMIT=1B"}, // NOTE: Absurdly low on purpose
 	}, func(config *docker.HostConfig) {

internal/datastore/crdb/crdb.go

@@ -29,6 +29,7 @@ import (
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	"github.com/authzed/spicedb/internal/datastore/revisions"
 	log "github.com/authzed/spicedb/internal/logging"
+	internalmetrics "github.com/authzed/spicedb/internal/metrics"
 	"github.com/authzed/spicedb/internal/sharederrors"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
@@ -72,6 +73,8 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 	}
 
 	includeQueryParametersInTraces := config.includeQueryParametersInTraces
+	pool.SetMetricsFactory(internalmetrics.NewPrometheusFactory(config.prometheusRegisterer))
+
 	readPoolConfig, err := pgxpool.ParseConfig(url)
 	if err != nil {
 		return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url)
@@ -213,7 +216,7 @@ func newCRDBDatastore(ctx context.Context, url string, options ...Option) (datas
 		return nil, common.RedactAndLogSensitiveConnString(ctx, errUnableToInstantiate, err, url)
 	}
 
-	err = ds.registerPrometheusCollectors(config.enablePrometheusStats)
+	err = ds.registerPrometheusCollectors(config.enablePrometheusStats, config.prometheusRegisterer)
 	if err != nil {
 		ds.cancel()
 		return nil, err
@@ -264,6 +267,7 @@ type crdbDatastore struct {
 	dburl                        string
 	readPool, writePool          *pool.RetryPool
 	collectors                   []prometheus.Collector
+	prometheusRegisterer         prometheus.Registerer
 	watchBufferLength            uint16
 	watchChangeBufferMaximumSize uint64
 	watchBufferWriteTimeout      time.Duration
@@ -481,9 +485,12 @@ func (cds *crdbDatastore) Close() error {
 	}
 	cds.readPool.Close()
 	cds.writePool.Close()
+	r := cds.prometheusRegisterer
+	if r == nil {
+		r = prometheus.DefaultRegisterer
+	}
 	for _, collector := range cds.collectors {
-		ok := prometheus.Unregister(collector)
-		if !ok {
+		if ok := r.Unregister(collector); !ok {
 			errs = append(errs, errors.New("could not unregister collector for CRDB datastore"))
 		}
 	}
@@ -658,7 +665,11 @@ func readClusterTTLNanos(ctx context.Context, conn pgxcommon.DBFuncQuerier) (int
 	return gcSeconds * 1_000_000_000, nil
 }
 
-func (cds *crdbDatastore) registerPrometheusCollectors(enablePrometheusStats bool) error {
+func (cds *crdbDatastore) registerPrometheusCollectors(enablePrometheusStats bool, registerer prometheus.Registerer) error {
+	if registerer == nil {
+		registerer = prometheus.DefaultRegisterer
+	}
+	cds.prometheusRegisterer = registerer
 	if !enablePrometheusStats {
 		return nil
 	}
@@ -667,21 +678,21 @@ func (cds *crdbDatastore) registerPrometheusCollectors(enablePrometheusStats boo
 		"db_name":    "spicedb",
 		"pool_usage": "read",
 	})
-
-	if err := prometheus.Register(readCollector); err != nil {
+	registered, err := internalmetrics.RegisterOrReuse(registerer, readCollector)
+	if err != nil {
 		return fmt.Errorf("failed to register prometheus read collector: %w", err)
 	}
-	cds.collectors = append(cds.collectors, readCollector)
+	cds.collectors = append(cds.collectors, registered)
 
 	writeCollector := pgxpoolprometheus.NewCollector(cds.writePool, map[string]string{
 		"db_name":    "spicedb",
 		"pool_usage": "write",
 	})
-
-	if err := prometheus.Register(writeCollector); err != nil {
+	registeredWrite, err := internalmetrics.RegisterOrReuse(registerer, writeCollector)
+	if err != nil {
 		return fmt.Errorf("failed to register prometheus write collector: %w", err)
 	}
-	cds.collectors = append(cds.collectors, writeCollector)
+	cds.collectors = append(cds.collectors, registeredWrite)
 
 	return nil
 }

internal/datastore/crdb/crdb_test.go

@@ -913,12 +913,12 @@ func TestRegisterPrometheusCollectors(t *testing.T) {
 		_ = cds.Close()
 	})
 
-	err = cds.registerPrometheusCollectors(false)
+	err = cds.registerPrometheusCollectors(false, nil)
 	require.NoError(t, err)
 	require.Empty(t, cds.collectors)
 
 	// Register collectors and verify the values of the metrics
-	err = cds.registerPrometheusCollectors(true)
+	err = cds.registerPrometheusCollectors(true, nil)
 	require.NoError(t, err)
 	require.Len(t, cds.collectors, 2)
 

internal/datastore/crdb/options.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
 	"k8s.io/utils/ptr"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
@@ -30,6 +31,7 @@ type crdbOptions struct {
 	analyzeBeforeStatistics        bool
 	filterMaximumIDCount           uint16
 	enablePrometheusStats          bool
+	prometheusRegisterer           prometheus.Registerer
 	withIntegrity                  bool
 	allowedMigrations              []string
 	columnOptimizationOption       common.ColumnOptimizationOption
@@ -345,6 +347,13 @@ func WithEnablePrometheusStats(enablePrometheusStats bool) Option {
 	return func(po *crdbOptions) { po.enablePrometheusStats = enablePrometheusStats }
 }
 
+// WithPrometheusRegisterer sets the prometheus.Registerer used to register
+// datastore metrics.  When not set (or nil), prometheus.DefaultRegisterer is
+// used.
+func WithPrometheusRegisterer(r prometheus.Registerer) Option {
+	return func(po *crdbOptions) { po.prometheusRegisterer = r }
+}
+
 // WithEnableConnectionBalancing marks whether Prometheus metrics provided by the Postgres
 // clients being used by the datastore are enabled.
 //

internal/datastore/crdb/pool/balancer.go

@@ -13,29 +13,40 @@ import (
 	"github.com/ccoveille/go-safecast/v2"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
-	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/sync/semaphore"
 
 	log "github.com/authzed/spicedb/internal/logging"
+	internalmetrics "github.com/authzed/spicedb/internal/metrics"
 	"github.com/authzed/spicedb/pkg/genutil"
 )
 
 var (
-	connectionsPerCRDBNodeCountGauge = prometheus.NewGaugeVec(prometheus.GaugeOpts{
+	connectionsPerCRDBNodeCountGauge internalmetrics.GaugeVec
+	pruningTimeHistogram             internalmetrics.HistogramVec
+)
+
+var (
+	connectionsGaugeOpts = internalmetrics.Opts{
 		Name: "crdb_connections_per_node",
 		Help: "The number of active connections SpiceDB holds to each CockroachDB node, by pool (read/write). Imbalanced values across nodes suggest the connection balancer is unable to redistribute connections evenly.",
-	}, []string{"pool", "node_id"})
-
-	pruningTimeHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
+	}
+	pruningHistogramOpts = internalmetrics.Opts{
 		Name:    "crdb_pruning_duration",
 		Help:    "Duration in milliseconds of one iteration of the CockroachDB connection balancer pruning excess connections from over-represented nodes. Elevated values indicate the balancer is struggling to rebalance connections.",
 		Buckets: []float64{.1, .2, .5, 1, 2, 5, 10, 20, 50, 100},
-	}, []string{"pool"})
+	}
 )
 
 func init() {
-	prometheus.MustRegister(connectionsPerCRDBNodeCountGauge)
-	prometheus.MustRegister(pruningTimeHistogram)
+	setBalancerMetricsFactory(internalmetrics.NewPrometheusFactory(nil))
+}
+
+func setBalancerMetricsFactory(factory internalmetrics.Factory) {
+	if factory == nil {
+		factory = internalmetrics.NoopFactory{}
+	}
+	connectionsPerCRDBNodeCountGauge = factory.GaugeVec(connectionsGaugeOpts, []string{"pool", "node_id"})
+	pruningTimeHistogram = factory.HistogramVec(pruningHistogramOpts, []string{"pool"})
 }
 
 type balancePoolConn[C balanceConn] interface {
@@ -182,10 +193,12 @@ func (p *nodeConnectionBalancer[P, C]) mustPruneConnections(ctx context.Context)
 	p.healthTracker.RLock()
 	for node := range p.healthTracker.nodesEverSeen {
 		if _, ok := connectionCounts[node]; !ok {
-			connectionsPerCRDBNodeCountGauge.DeletePartialMatch(map[string]string{
-				"pool":    p.pool.ID(),
-				"node_id": strconv.FormatUint(uint64(node), 10),
-			})
+			if pgv, ok := internalmetrics.AsPrometheusGaugeVec(connectionsPerCRDBNodeCountGauge); ok {
+				pgv.DeletePartialMatch(map[string]string{
+					"pool":    p.pool.ID(),
+					"node_id": strconv.FormatUint(uint64(node), 10),
+				})
+			}
 		}
 	}
 	p.healthTracker.RUnlock()

internal/datastore/crdb/pool/health.go

@@ -7,24 +7,33 @@ import (
 	"time"
 
 	"github.com/jackc/pgx/v5"
-	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/time/rate"
 
 	"github.com/authzed/jitterbug"
 
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	log "github.com/authzed/spicedb/internal/logging"
+	internalmetrics "github.com/authzed/spicedb/internal/metrics"
 )
 
 const errorBurst = 2
 
-var healthyCRDBNodeCountGauge = prometheus.NewGauge(prometheus.GaugeOpts{
+var healthyCRDBNodeCountGauge internalmetrics.Gauge
+
+var healthyNodesGaugeOpts = internalmetrics.Opts{
 	Name: "crdb_healthy_nodes",
 	Help: "the number of healthy crdb nodes detected by spicedb",
-})
+}
 
 func init() {
-	prometheus.MustRegister(healthyCRDBNodeCountGauge)
+	setHealthMetricsFactory(internalmetrics.NewPrometheusFactory(nil))
+}
+
+func setHealthMetricsFactory(factory internalmetrics.Factory) {
+	if factory == nil {
+		factory = internalmetrics.NoopFactory{}
+	}
+	healthyCRDBNodeCountGauge = factory.Gauge(healthyNodesGaugeOpts)
 }
 
 // NodeHealthTracker detects changes in the node pool by polling the cluster periodically and recording

internal/datastore/crdb/pool/metrics.go

@@ -0,0 +1,10 @@
+package pool
+
+import internalmetrics "github.com/authzed/spicedb/internal/metrics"
+
+// SetMetricsFactory configures which metrics factory is used by this package.
+func SetMetricsFactory(factory internalmetrics.Factory) {
+	setBalancerMetricsFactory(factory)
+	setHealthMetricsFactory(factory)
+	setPoolMetricsFactory(factory)
+}

internal/datastore/crdb/pool/metrics_test.go

@@ -0,0 +1,47 @@
+package pool
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	internalmetrics "github.com/authzed/spicedb/internal/metrics"
+)
+
+func TestSetMetricsFactoryAppliesToAllPoolMetrics(t *testing.T) {
+	recording := internalmetrics.NewRecordingFactory()
+	SetMetricsFactory(recording)
+	t.Cleanup(func() {
+		SetMetricsFactory(internalmetrics.NewPrometheusFactory(nil))
+	})
+
+	resetHistogram.Observe(1)
+	healthyCRDBNodeCountGauge.Set(2)
+	connectionsPerCRDBNodeCountGauge.WithLabelValues("read", "1").Set(3)
+	pruningTimeHistogram.WithLabelValues("read").Observe(4)
+
+	require.Equal(t, uint64(1), recording.HistogramCount("", "", "crdb_client_resets"))
+	require.InEpsilon(t, 2.0, recording.GaugeValue("", "", "crdb_healthy_nodes"), 1e-9)
+	require.InEpsilon(t, 3.0, recording.GaugeVecValue("", "", "crdb_connections_per_node", "read", "1"), 1e-9)
+	require.Equal(t, uint64(1), recording.HistogramVecCount("", "", "crdb_pruning_duration", "read"))
+}
+
+func TestSetMetricsFactoryWithNilAndNoopPool(t *testing.T) {
+	SetMetricsFactory(internalmetrics.NoopFactory{})
+	require.NotPanics(t, func() {
+		resetHistogram.Observe(1)
+		healthyCRDBNodeCountGauge.Set(1)
+		connectionsPerCRDBNodeCountGauge.WithLabelValues("read", "1").Set(1)
+		pruningTimeHistogram.WithLabelValues("read").Observe(1)
+	})
+
+	SetMetricsFactory(nil)
+	require.NotPanics(t, func() {
+		resetHistogram.Observe(1)
+		healthyCRDBNodeCountGauge.Set(1)
+		connectionsPerCRDBNodeCountGauge.WithLabelValues("read", "1").Set(1)
+		pruningTimeHistogram.WithLabelValues("read").Observe(1)
+	})
+
+	SetMetricsFactory(internalmetrics.NewPrometheusFactory(nil))
+}

internal/datastore/crdb/pool/pool.go

@@ -11,11 +11,11 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
-	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/time/rate"
 
 	"github.com/authzed/spicedb/internal/datastore/postgres/common"
 	log "github.com/authzed/spicedb/internal/logging"
+	internalmetrics "github.com/authzed/spicedb/internal/metrics"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
@@ -28,14 +28,23 @@ type pgxPool interface {
 	Stat() *pgxpool.Stat
 }
 
-var resetHistogram = prometheus.NewHistogram(prometheus.HistogramOpts{
+var resetHistogram internalmetrics.Histogram
+
+var resetHistogramOpts = internalmetrics.Opts{
 	Name:    "crdb_client_resets",
 	Help:    "Distribution of the number of client-side transaction restarts per transaction attempt. Restarts occur when CockroachDB returns a serialization failure (40001) and the driver retries the transaction from scratch. Sustained high values indicate transaction contention.",
 	Buckets: []float64{0, 1, 2, 5, 10, 20, 50},
-})
+}
 
 func init() {
-	prometheus.MustRegister(resetHistogram)
+	setPoolMetricsFactory(internalmetrics.NewPrometheusFactory(nil))
+}
+
+func setPoolMetricsFactory(factory internalmetrics.Factory) {
+	if factory == nil {
+		factory = internalmetrics.NoopFactory{}
+	}
+	resetHistogram = factory.Histogram(resetHistogramOpts)
 }
 
 type ctxDisableRetries struct{}