uniugm/admin/lib/prom/promethus.go

package prom

import (
	"fmt"
	"sync/atomic"

	ginPprof "github.com/gin-contrib/pprof"
	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

func NewCounter(name string) *PromeCounterStatMgr {
	mgr := new(PromeCounterStatMgr)
	mgr.promeStatMgr = newPromeStatMgr(name, func(defaultLabels, labels []string) prometheus.Collector {
		return newCounterVec(name, defaultLabels, labels)
	})
	return mgr
}

func NewGauge(name string) *PromeGaugeStatMgr {
	mgr := new(PromeGaugeStatMgr)
	mgr.promeStatMgr = newPromeStatMgr(name, func(defaultLabels, labels []string) prometheus.Collector {
		return newGagueVec(name, defaultLabels, labels)
	})
	return mgr
}

func NewHistogram(name string, buckets []float64) *PromeHistogramStatMgr {
	mgr := new(PromeHistogramStatMgr)
	mgr.promeStatMgr = newPromeStatMgr(name, func(defaultLabels, labels []string) prometheus.Collector {
		return newHistogramVec(name, defaultLabels, labels, buckets)
	})
	return mgr
}

type PromeCounterStatMgr struct {
	*promeStatMgr
}

// InitLabels 初始化指标的标签说明，以后每次收集指标都给上标签值
func (c *PromeCounterStatMgr) InitLabels(labels []string) *PromeCounterStatMgr {
	return c.InitDefaultLabels(nil, labels)
}

// InitDefaultLabels 带上默认标签，例如某些服务启动就能知道标签值，比如app_name，app_id等，
// 又不想每次收集指标都传递标签值，可以用此方法一开始就设置好
func (c *PromeCounterStatMgr) InitDefaultLabels(defaultLabels map[string]string, labels []string) *PromeCounterStatMgr {
	c.promeStatMgr.withDefaultLabels(defaultLabels, labels)
	return c
}

// LabelValues 准备收集指标，调用这个传递标签值，后续再给上指标值
// 即：c.LabelValues("xxx").Add(234)
func (c *PromeCounterStatMgr) LabelValues(labels ...string) prometheus.Counter {
	return c.promeStatMgr.getCounterWithLabels(labels...)
}

type PromeGaugeStatMgr struct {
	*promeStatMgr
}

// InitLabels 初始化指标的标签说明，以后每次收集指标都给上标签值
func (c *PromeGaugeStatMgr) InitLabels(labels []string) *PromeGaugeStatMgr {
	return c.InitDefaultLabels(nil, labels)
}

// InitDefaultLabels 带上默认标签，例如某些服务启动就能知道标签值，比如app_name，app_id等，
// 又不想每次收集指标都传递标签值，可以用此方法一开始就设置好
func (c *PromeGaugeStatMgr) InitDefaultLabels(defaultLabels map[string]string, labels []string) *PromeGaugeStatMgr {
	c.promeStatMgr.withDefaultLabels(defaultLabels, labels)
	return c
}

// LabelValues 准备收集指标，调用这个传递标签值，后续再给上指标值
// 即：c.LabelValues("xxx").Add(234)
func (c *PromeGaugeStatMgr) LabelValues(labels ...string) prometheus.Gauge {
	return c.promeStatMgr.getGaugeWithLabels(labels...)
}

type PromeHistogramStatMgr struct {
	*promeStatMgr
}

// InitLabels 初始化指标的标签说明，以后每次收集指标都给上标签值
func (c *PromeHistogramStatMgr) InitLabels(labels []string) *PromeHistogramStatMgr {
	return c.InitDefaultLabels(nil, labels)
}

// InitDefaultLabels 带上默认标签，例如某些服务启动就能知道标签值，比如app_name，app_id等，
// 又不想每次收集指标都传递标签值，可以用此方法一开始就设置好
func (c *PromeHistogramStatMgr) InitDefaultLabels(defaultLabels map[string]string, labels []string) *PromeHistogramStatMgr {
	c.promeStatMgr.withDefaultLabels(defaultLabels, labels)
	return c
}

// LabelValues 准备收集指标，调用这个传递标签值，后续再给上指标值
// 即：c.LabelValues("xxx").Add(234)
func (c *PromeHistogramStatMgr) LabelValues(labels ...string) prometheus.Observer {
	return c.promeStatMgr.getHistogramWithLabels(labels...)
}

type promeStatMgr struct {
	collector          prometheus.Collector
	name               string
	newCollectorFun    func(defaultLabels, labels []string) prometheus.Collector
	defaultLabelsValue []string
	initCounter        int32
}

func (mgr *promeStatMgr) withDefaultLabels(defaultLabels map[string]string, labels []string) {
	if atomic.AddInt32(&mgr.initCounter, 1) > 1 {
		panic(fmt.Errorf("promethus vec labels must init once"))
	}

	defaultLabelKeys := make([]string, 0)
	defaultLabelValues := make([]string, 0)
	for k, v := range defaultLabels {
		defaultLabelKeys = append(defaultLabelKeys, k)
		defaultLabelValues = append(defaultLabelValues, v)
	}

	mgr.defaultLabelsValue = defaultLabelValues
	mgr.collector = mgr.newCollectorFun(defaultLabelKeys, labels)
	return
}

func (mgr *promeStatMgr) joinLabels(labels ...string) []string {
	newLabels := make([]string, 0, len(mgr.defaultLabelsValue)+len(labels))
	newLabels = append(newLabels, mgr.defaultLabelsValue...)
	newLabels = append(newLabels, labels...)
	return newLabels
}

func (mgr *promeStatMgr) getCounterWithLabels(labels ...string) prometheus.Counter {
	newLabels := mgr.joinLabels(labels...)
	return mgr.collector.(*prometheus.CounterVec).WithLabelValues(newLabels...)
}

func (mgr *promeStatMgr) getGaugeWithLabels(labels ...string) prometheus.Gauge {
	newLabels := mgr.joinLabels(labels...)
	return mgr.collector.(*prometheus.GaugeVec).WithLabelValues(newLabels...)
}

func (mgr *promeStatMgr) getHistogramWithLabels(labels ...string) prometheus.Observer {
	newLabels := mgr.joinLabels(labels...)
	return mgr.collector.(*prometheus.HistogramVec).WithLabelValues(newLabels...)
}

func newPromeStatMgr(name string,
	newCollectorFun func(defaultLabelsKey []string, labels []string) prometheus.Collector) *promeStatMgr {
	mgr := new(promeStatMgr)
	mgr.name = name
	mgr.newCollectorFun = newCollectorFun
	return mgr
}

func newCounterVec(name string, defaultLabels []string, dynamicLabels []string) prometheus.Collector {
	vec := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: name,
			Help: name,
		},
		append(defaultLabels, dynamicLabels...),
	)
	prometheus.MustRegister(vec)
	return vec
}

func newGagueVec(name string, defaultLabels []string, dynamicLabels []string) prometheus.Collector {
	vec := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: name,
			Help: name,
		},
		append(defaultLabels, dynamicLabels...),
	)
	prometheus.MustRegister(vec)
	return vec
}

func newHistogramVec(name string, defaultLabels []string, dynamicLabels []string, buckets []float64) prometheus.Collector {
	vec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    name,
			Help:    name,
			Buckets: buckets,
		},
		append(defaultLabels, dynamicLabels...),
	)
	prometheus.MustRegister(vec)
	return vec
}

func NewEngine(enablePprof bool) *gin.Engine {
	//gin.SetMode(gin.ReleaseMode)
	engine := gin.New()
	engine.Use(gin.Recovery())
	RouteEngine(engine, enablePprof)
	return engine
}

func RouteEngine(engine *gin.Engine, enablePprof bool) {
	if enablePprof {
		ginPprof.Register(engine)
	}
	engine.Use(gin.LoggerWithConfig(gin.LoggerConfig{
		SkipPaths: []string{"/metrics"},
	}))
	engine.GET("/metrics", gin.WrapH(promhttp.Handler()))
}