Tally based RED metrics middleware

redMetricsMiddleware.go

package middleware

import (
	"net/http"

	"github.com/uber-go/tally/v4"
)

type statusCodeAwareWriter struct {
	http.ResponseWriter
	statusCode int
}

func (w *statusCodeAwareWriter) WriteHeader(status int) {
	w.statusCode = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusCodeAwareWriter) Write(b []byte) (int, error) {
	return w.ResponseWriter.Write(b)
}

func (w *statusCodeAwareWriter) HasAnError() bool {
	return w.statusCode >= http.StatusInternalServerError
}

type RedMetricsMiddleware struct {
	TotalRequests    tally.Counter
	RequestsDuration tally.Histogram
	RequestFailures  tally.Counter
}

func NewRedMetricsMiddleware(c tally.Counter, d tally.Histogram, f tally.Counter) *RedMetricsMiddleware {
	return &RedMetricsMiddleware{
		TotalRequests:    c,
		RequestsDuration: d,
		RequestFailures:  f,
	}
}

func (m *RedMetricsMiddleware) Handle() func(next http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			dw := &statusCodeAwareWriter{ResponseWriter: w}
			stopwatch := m.RequestsDuration.Start()
			m.TotalRequests.Inc(1)
			next.ServeHTTP(dw, r)
			stopwatch.Stop()
			if dw.HasAnError() {
				m.RequestFailures.Inc(1)
			}
		}
		return http.HandlerFunc(fn)
	}
}