fix: protect publish path from mysql exhaustion

server/metrics/db_pool.go

@@ -0,0 +1,80 @@
+package metrics
+
+import (
+	"database/sql"
+	"strings"
+	"sync"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+type dbStatsProvider interface {
+	Stats() sql.DBStats
+}
+
+type dbPoolCollector struct {
+	pools sync.Map
+	desc  map[string]*prometheus.Desc
+}
+
+func newDBPoolCollector() *dbPoolCollector {
+	labels := []string{"pod", "pool"}
+	return &dbPoolCollector{desc: map[string]*prometheus.Desc{
+		"max_open":            prometheus.NewDesc("im_db_pool_max_open_connections", "Configured maximum open database connections.", labels, nil),
+		"open":                prometheus.NewDesc("im_db_pool_open_connections", "Current open database connections.", labels, nil),
+		"in_use":              prometheus.NewDesc("im_db_pool_in_use_connections", "Current database connections in use.", labels, nil),
+		"idle":                prometheus.NewDesc("im_db_pool_idle_connections", "Current idle database connections.", labels, nil),
+		"wait":                prometheus.NewDesc("im_db_pool_wait_total", "Total waits for a database connection.", labels, nil),
+		"wait_duration":       prometheus.NewDesc("im_db_pool_wait_duration_seconds_total", "Total time spent waiting for database connections.", labels, nil),
+		"max_idle_closed":     prometheus.NewDesc("im_db_pool_max_idle_closed_total", "Connections closed because the idle limit was exceeded.", labels, nil),
+		"max_lifetime_closed": prometheus.NewDesc("im_db_pool_max_lifetime_closed_total", "Connections closed because their lifetime expired.", labels, nil),
+	}}
+}
+
+func (c *dbPoolCollector) Describe(ch chan<- *prometheus.Desc) {
+	for _, desc := range c.desc {
+		ch <- desc
+	}
+}
+
+func (c *dbPoolCollector) Collect(ch chan<- prometheus.Metric) {
+	pod := metricPodName()
+	c.pools.Range(func(key, value any) bool {
+		pool, ok := key.(string)
+		provider, providerOK := value.(dbStatsProvider)
+		if !ok || !providerOK {
+			return true
+		}
+		stats := provider.Stats()
+		labels := []string{pod, pool}
+		ch <- prometheus.MustNewConstMetric(c.desc["max_open"], prometheus.GaugeValue, float64(stats.MaxOpenConnections), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["open"], prometheus.GaugeValue, float64(stats.OpenConnections), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["in_use"], prometheus.GaugeValue, float64(stats.InUse), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["idle"], prometheus.GaugeValue, float64(stats.Idle), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["wait"], prometheus.CounterValue, float64(stats.WaitCount), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["wait_duration"], prometheus.CounterValue, stats.WaitDuration.Seconds(), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["max_idle_closed"], prometheus.CounterValue, float64(stats.MaxIdleClosed), labels...)
+		ch <- prometheus.MustNewConstMetric(c.desc["max_lifetime_closed"], prometheus.CounterValue, float64(stats.MaxLifetimeClosed), labels...)
+		return true
+	})
+}
+
+var dbPools = newDBPoolCollector()
+
+// RegisterDBPool exposes database/sql pool pressure for one logical database.
+// Registering the same name again replaces the provider, which keeps tests and
+// adapter reopen operations safe.
+func RegisterDBPool(name string, provider dbStatsProvider) {
+	if provider == nil {
+		return
+	}
+	name = strings.TrimSpace(name)
+	if name == "" {
+		name = "default"
+	}
+	dbPools.pools.Store(name, provider)
+}
+
+func init() {
+	Registry.MustRegister(dbPools)
+}

server/metrics/db_pool_test.go

@@ -0,0 +1,47 @@
+package metrics
+
+import (
+	"database/sql"
+	"testing"
+	"time"
+)
+
+type fixedDBStats struct {
+	stats sql.DBStats
+}
+
+func (s fixedDBStats) Stats() sql.DBStats { return s.stats }
+
+func TestDBPoolCollector(t *testing.T) {
+	RegisterDBPool("test", fixedDBStats{stats: sql.DBStats{
+		MaxOpenConnections: 32,
+		OpenConnections:    12,
+		InUse:              7,
+		Idle:               5,
+		WaitCount:          9,
+		WaitDuration:       1500 * time.Millisecond,
+	}})
+
+	families, err := Registry.Gather()
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := map[string]bool{
+		"im_db_pool_max_open_connections":        false,
+		"im_db_pool_open_connections":            false,
+		"im_db_pool_in_use_connections":          false,
+		"im_db_pool_idle_connections":            false,
+		"im_db_pool_wait_total":                  false,
+		"im_db_pool_wait_duration_seconds_total": false,
+	}
+	for _, family := range families {
+		if _, ok := want[family.GetName()]; ok {
+			want[family.GetName()] = true
+		}
+	}
+	for name, found := range want {
+		if !found {
+			t.Errorf("metric %s was not gathered", name)
+		}
+	}
+}