byop-engine/middleware/metrics.go

package middleware

import (
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	// RequestsTotal counts the number of HTTP requests processed
	RequestsTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "http_requests_total",
			Help: "Total number of HTTP requests",
		},
		[]string{"method", "path", "status"},
	)

	// RequestDuration observes the HTTP request duration
	RequestDuration = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "http_request_duration_seconds",
			Help:    "HTTP request duration in seconds",
			Buckets: prometheus.DefBuckets,
		},
		[]string{"method", "path"},
	)

	// ResponseSize observes the HTTP response size
	ResponseSize = promauto.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "http_response_size_bytes",
			Help:    "HTTP response size in bytes",
			Buckets: prometheus.ExponentialBuckets(100, 10, 8), // From 100B to 10GB
		},
		[]string{"method", "path"},
	)
)

// Metrics is a middleware that collects Prometheus metrics for HTTP requests
func Metrics(c *gin.Context) {
	start := time.Now()

	// Create a custom response writer to capture status code and response size
	mrw := &metricsResponseWriter{ResponseWriter: c.Writer}
	c.Writer = mrw

	// Process the request
	c.Next()

	// Record metrics
	duration := time.Since(start).Seconds()
	statusCode := strconv.Itoa(mrw.statusCode)
	method := c.Request.Method
	path := c.Request.URL.Path

	RequestsTotal.WithLabelValues(method, path, statusCode).Inc()
	RequestDuration.WithLabelValues(method, path).Observe(duration)
	ResponseSize.WithLabelValues(method, path).Observe(float64(mrw.responseSize))
}

// metricsResponseWriter is a custom ResponseWriter that captures status code and response size
type metricsResponseWriter struct {
	gin.ResponseWriter
	statusCode   int
	responseSize int
}

// WriteHeader captures the status code
func (mrw *metricsResponseWriter) WriteHeader(code int) {
	mrw.statusCode = code
	mrw.ResponseWriter.WriteHeader(code)
}

// Write captures the response size
func (mrw *metricsResponseWriter) Write(b []byte) (int, error) {
	n, err := mrw.ResponseWriter.Write(b)
	if err != nil {
		return n, err
	}
	mrw.responseSize += n
	return n, nil
}