gitlab.com/gitlab-org/labkit@v1.21.0/metrics/sqlmetrics/dbstats.go

package sqlmetrics

import (
	"database/sql"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	namespace   = "go_sql_dbstats"
	subsystem   = "connections"
	dbNameLabel = "db_name"

	// Names for the recorded metrics.
	maxOpenConnectionsName = "max_open"
	openConnectionsName    = "open"
	inUseName              = "in_use"
	idleName               = "idle"
	waitCountName          = "waits_total"
	waitDurationName       = "wait_seconds_total"
	maxIdleClosedName      = "max_idle_closed_count_total"
	maxIdleTimeClosedName  = "max_idle_time_closed_count_total"
	maxLifetimeClosedName  = "max_lifetime_closed_count_total"

	// Descriptions for the recorded metrics.
	maxOpenConnectionsDesc = "The limit of open connections to the database."
	openConnectionsDesc    = "The number of established connections both in use and idle."
	inUseDesc              = "The number of connections currently in use."
	idleDesc               = "The number of idle connections."
	waitCountDesc          = "The total number of connections waited for."
	waitDurationDesc       = "The total time blocked waiting for a new connection."
	maxIdleClosedDesc      = "The total number of connections closed due to SetMaxIdleConns."
	maxIdleTimeClosedDesc  = "The total number of connections closed due to SetConnMaxIdleTime."
	maxLifetimeClosedDesc  = "The total number of connections closed due to SetConnMaxLifetime."
)

// DBStatsGetter is an interface for sql.DBStats. It's implemented by sql.DB.
type DBStatsGetter interface {
	Stats() sql.DBStats
}

// DBStatsCollector implements the prometheus.Collector interface.
type DBStatsCollector struct {
	sg DBStatsGetter

	maxOpenDesc           *prometheus.Desc
	openDesc              *prometheus.Desc
	inUseDesc             *prometheus.Desc
	idleDesc              *prometheus.Desc
	waitCountDesc         *prometheus.Desc
	waitDurationDesc      *prometheus.Desc
	maxIdleClosedDesc     *prometheus.Desc
	maxIdleTimeClosedDesc *prometheus.Desc
	maxLifetimeClosedDesc *prometheus.Desc
}

// Describe implements the prometheus.Collector interface.
func (c *DBStatsCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.maxOpenDesc
	ch <- c.openDesc
	ch <- c.inUseDesc
	ch <- c.idleDesc
	ch <- c.waitCountDesc
	ch <- c.waitDurationDesc
	ch <- c.maxIdleClosedDesc
	ch <- c.maxIdleTimeClosedDesc
	ch <- c.maxLifetimeClosedDesc
}

// Collect implements the prometheus.Collector interface.
func (c *DBStatsCollector) Collect(ch chan<- prometheus.Metric) {
	stats := c.sg.Stats()

	ch <- prometheus.MustNewConstMetric(
		c.maxOpenDesc,
		prometheus.GaugeValue,
		float64(stats.MaxOpenConnections),
	)
	ch <- prometheus.MustNewConstMetric(
		c.openDesc,
		prometheus.GaugeValue,
		float64(stats.OpenConnections),
	)
	ch <- prometheus.MustNewConstMetric(
		c.inUseDesc,
		prometheus.GaugeValue,
		float64(stats.InUse),
	)
	ch <- prometheus.MustNewConstMetric(
		c.idleDesc,
		prometheus.GaugeValue,
		float64(stats.Idle),
	)
	ch <- prometheus.MustNewConstMetric(
		c.waitCountDesc,
		prometheus.CounterValue,
		float64(stats.WaitCount),
	)
	ch <- prometheus.MustNewConstMetric(
		c.waitDurationDesc,
		prometheus.CounterValue,
		stats.WaitDuration.Seconds(),
	)
	ch <- prometheus.MustNewConstMetric(
		c.maxIdleClosedDesc,
		prometheus.CounterValue,
		float64(stats.MaxIdleClosed),
	)
	ch <- prometheus.MustNewConstMetric(
		c.maxIdleTimeClosedDesc,
		prometheus.CounterValue,
		float64(stats.MaxIdleTimeClosed),
	)
	ch <- prometheus.MustNewConstMetric(
		c.maxLifetimeClosedDesc,
		prometheus.CounterValue,
		float64(stats.MaxLifetimeClosed),
	)
}

type dbStatsCollectorConfig struct {
	extraLabels prometheus.Labels
}

// DBStatsCollectorOption is used to pass options in NewDBStatsCollector.
type DBStatsCollectorOption func(*dbStatsCollectorConfig)

// WithExtraLabels will configure extra label values to apply to the DBStats metrics.
// A label named db_name will be ignored, as this is set internally.
func WithExtraLabels(labelValues map[string]string) DBStatsCollectorOption {
	return func(config *dbStatsCollectorConfig) {
		config.extraLabels = labelValues
	}
}

func applyDBStatsCollectorOptions(opts []DBStatsCollectorOption) dbStatsCollectorConfig {
	config := dbStatsCollectorConfig{}
	for _, v := range opts {
		v(&config)
	}

	return config
}

// NewDBStatsCollector creates a new DBStatsCollector.
func NewDBStatsCollector(dbName string, sg DBStatsGetter, opts ...DBStatsCollectorOption) *DBStatsCollector {
	config := applyDBStatsCollectorOptions(opts)

	if config.extraLabels == nil {
		config.extraLabels = make(prometheus.Labels)
	}
	config.extraLabels[dbNameLabel] = dbName

	return &DBStatsCollector{
		sg: sg,
		maxOpenDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, maxOpenConnectionsName),
			maxOpenConnectionsDesc,
			nil,
			config.extraLabels,
		),
		openDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, openConnectionsName),
			openConnectionsDesc,
			nil,
			config.extraLabels,
		),
		inUseDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, inUseName),
			inUseDesc,
			nil,
			config.extraLabels,
		),
		idleDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, idleName),
			idleDesc,
			nil,
			config.extraLabels,
		),
		waitCountDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, waitCountName),
			waitCountDesc,
			nil,
			config.extraLabels,
		),
		waitDurationDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, waitDurationName),
			waitDurationDesc,
			nil,
			config.extraLabels,
		),
		maxIdleClosedDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, maxIdleClosedName),
			maxIdleClosedDesc,
			nil,
			config.extraLabels,
		),
		maxIdleTimeClosedDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, maxIdleTimeClosedName),
			maxIdleTimeClosedDesc,
			nil,
			config.extraLabels,
		),
		maxLifetimeClosedDesc: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, subsystem, maxLifetimeClosedName),
			maxLifetimeClosedDesc,
			nil,
			config.extraLabels,
		),
	}
}