fetch_ml/internal/api/metrics_middleware.go

63 lines
1.4 KiB
Go

package api
import (
"bufio"
"fmt"
"net"
"net/http"
"time"
)
// wrapWithMetrics wraps a handler with Prometheus metrics tracking
func (s *Server) wrapWithMetrics(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if s.promMetrics == nil {
next.ServeHTTP(w, r)
return
}
start := time.Now()
// Track HTTP request
method := r.Method
endpoint := r.URL.Path
// Wrap response writer to capture status code
ww := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
// Serve request
next.ServeHTTP(ww, r)
// Record metrics
duration := time.Since(start)
statusStr := http.StatusText(ww.statusCode)
s.promMetrics.IncHTTPRequests(method, endpoint, statusStr)
s.promMetrics.ObserveHTTPDuration(method, endpoint, duration)
})
}
// responseWriter wraps http.ResponseWriter to capture status code
type responseWriter struct {
http.ResponseWriter
statusCode int
}
func (rw *responseWriter) WriteHeader(code int) {
rw.statusCode = code
rw.ResponseWriter.WriteHeader(code)
}
func (rw *responseWriter) Flush() {
if f, ok := rw.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
h, ok := rw.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("websocket: response does not implement http.Hijacker")
}
return h.Hijack()
}