kit/middleware/prome/prome.go

package prome

import (
	"context"
	"fmt"
	"git.diulo.com/mogfee/kit/middleware"
	"git.diulo.com/mogfee/kit/transport"
	"github.com/prometheus/client_golang/prometheus"
	"time"
)

const namespace = "service"

var (
	labels = []string{"status", "endpoint", "method"}

	uptime = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: namespace,
		Name:      "uptime",
		Help:      "HTTP service uptime.",
	}, nil)

	reqCount = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: namespace,
		Name:      "http_request_count_total",
		Help:      "Total number of HTTP requests made.",
	}, labels)

	reqDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace: namespace,
		Name:      "http_request_duration_seconds",
		Help:      "HTTP request latencies in seconds.",
	}, labels)
	reqSizeBytes = prometheus.NewSummaryVec(prometheus.SummaryOpts{
		Namespace: namespace,
		Name:      "http_request_size_bytes",
		Help:      "HTTP request sizes in bytes.",
	}, labels)
	respSizeBytes = prometheus.NewSummaryVec(prometheus.SummaryOpts{
		Namespace: namespace,
		Name:      "http_response_size_bytes",
		Help:      "HTTP request sizes in bytes.",
	}, labels)
)

func init() {
	prometheus.MustRegister(uptime, reqCount, reqDuration, reqSizeBytes, respSizeBytes)
	go recordUptime()
}
func recordUptime() {
	for range time.Tick(time.Second) {
		uptime.WithLabelValues().Inc()
	}
}

func Prometheus() middleware.Middleware {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, a any) (any, error) {
			tr, ok := transport.FromServerContext(ctx)
			if !ok {
				return handler(ctx, a)
			}
			start := time.Now()
			resp, err := handler(ctx, a)
			code := "500"
			if err != nil {
				switch aerr := err.(type) {
				case *errorx.Error:
					code = fmt.Sprintf("%d", aerr.Code)
				default:
					code = "500"
				}
			}
			lvs := []string{code, tr.Endpoint(), tr.Operation()}

			reqCount.WithLabelValues(lvs...).Inc()
			reqDuration.WithLabelValues(lvs...).Observe(time.Since(start).Seconds())

			return resp, err
		}
	}
}