sparrow/pkg/server/promethues.go

package server

import (
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"net/http"
)

var rpcCallCnt = &Metric{
	ID:          "rpcCallCnt",
	Name:        "rpc_call_total",
	Description: "HOW MANY RPC CALL",
	Type:        "counter",
}

var rpcCallDur = &Metric{
	ID:          "rpc_call_duration",
	Name:        "rpc_call_duration_seconds",
	Description: "rpc call duration seconds",
	Type:        "histogram_vec",
	Args:        []string{"client", "server", "method"},
}

var standardMetrics = []*Metric{
	rpcCallCnt,
	rpcCallDur,
}

type Metric struct {
	MetricCollector prometheus.Collector
	ID              string
	Name            string
	Description     string
	Type            string
	Args            []string
}

type Prometheus struct {
	MetricsList []*Metric
	MetricsPath string
	CallCnt     prometheus.Counter
	CallDur     *prometheus.HistogramVec
}

func NewPrometheus(subsystem string, customMetricsList ...*Metric) *Prometheus {
	var metricsList []*Metric
	for _, v := range customMetricsList {
		metricsList = append(metricsList, v)
	}
	for _, metric := range standardMetrics {
		metricsList = append(metricsList, metric)
	}
	p := &Prometheus{
		MetricsList: metricsList,
		MetricsPath: "/metrics",
	}
	p.registerMetrics(subsystem)
	return p
}

func (p *Prometheus) Start(addr string) error {
	http.Handle(p.MetricsPath, promhttp.Handler())
	return http.ListenAndServe(addr, nil)
}

func (p *Prometheus) RegisterMetrics(sub string, metrics ...*Metric) {
	for _, m := range metrics {
		metric := NewMetric(m, sub)
		if err := prometheus.Register(metric); err != nil {
			Log.Errorf("%s could not be registered in prometheus", m.Name)
		}
	}
}

// NewMetric associates prometheus.Collector based on Metric.Type
func NewMetric(m *Metric, subsystem string) prometheus.Collector {
	var metric prometheus.Collector
	switch m.Type {
	case "counter_vec":
		metric = prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "counter":
		metric = prometheus.NewCounter(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "gauge_vec":
		metric = prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "gauge":
		metric = prometheus.NewGauge(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "histogram_vec":
		metric = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "histogram":
		metric = prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "summary_vec":
		metric = prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "summary":
		metric = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	}
	return metric
}

func (p *Prometheus) registerMetrics(sub string) {
	for _, metricDef := range p.MetricsList {
		metric := NewMetric(metricDef, sub)
		if err := prometheus.Register(metric); err != nil {
			Log.Errorf("%s could not be registered in prometheus", metricDef.Name)
		}
		switch metricDef {
		case rpcCallCnt:
			p.CallCnt = metric.(prometheus.Counter)
		case rpcCallDur:
			p.CallDur = metric.(*prometheus.HistogramVec)
		}
		metricDef.MetricCollector = metric
	}
}