webhooker/internal/middleware/metrics_test.go

package middleware_test

import (
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/go-chi/chi"
	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	prommetrics "github.com/slok/go-http-metrics/metrics/prometheus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/webhooker/internal/config"
	"sneak.berlin/go/webhooker/internal/middleware"
)

const (
	// metricsProbePaths is how many distinct receiver paths each
	// cardinality assertion drives. The defect these tests pin cost
	// roughly 26 permanent series per distinct path, so a couple of
	// hundred puts a regression thousands of series over the bound
	// rather than leaving it to a rounding argument.
	metricsProbePaths = 250

	// receiverRoutePattern is the one handler label every receiver
	// request must produce, however the client varies the path.
	receiverRoutePattern = "/webhook/{uuid}"

	// okRoute is a static route used to pin that the response-writer
	// interceptor still reports status and size after the handler id
	// stopped coming from the URL.
	okRoute = "/ok"

	// okBody is what okRoute writes, so the recorded response size is
	// a number the test knows.
	okBody = "ok"

	// generousReceiverLimit is a per-entrypoint receiver limit high
	// enough that no probe in this file trips the limiter unless it
	// means to.
	generousReceiverLimit = 100000

	// tightReceiverLimit forces the receiver's aggregate limiter to
	// reject: the aggregate ceiling is ten times this, so a probe of
	// metricsProbePaths requests spends it many times over.
	tightReceiverLimit = 1
)

// metricsTestRouter builds a router whose middleware ordering mirrors
// the real server's: the metrics recorder is GLOBAL, installed by
// Server.setupGlobalMiddleware before chi has matched anything, and
// the receiver rate limiter is ROUTE-LEVEL, installed by
// Server.setupWebhookRoutes inside it. That ordering is the whole
// defect, so a test that flattens it would prove nothing.
//
// The recorder writes to a registry of the test's own rather than the
// process-wide default one, so each test observes only its own
// traffic.
func metricsTestRouter(
	t *testing.T,
	receiverLimit int,
) (http.Handler, *prometheus.Registry) {
	t.Helper()

	log := slog.New(slog.DiscardHandler)
	cfg := &config.Config{
		Environment:       "prod",
		ReceiverRateLimit: receiverLimit,
	}
	m := middleware.NewForTest(
		log, cfg, newTestSessionManager(cfg, log, nil),
	)

	reg := prometheus.NewRegistry()
	rec := prommetrics.NewRecorder(prommetrics.Config{Registry: reg})

	r := chi.NewRouter()
	r.Use(middleware.MetricsMiddlewareForTest(rec))

	r.Get(okRoute, func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(okBody))
	})

	// The real receiver answers 404 for a UUID naming no stored
	// entrypoint, which is what every invented path here is.
	r.With(m.ReceiverRateLimit()).HandleFunc(
		receiverRoutePattern,
		func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusNotFound)
		},
	)

	return r, reg
}

// probe is one request a cardinality assertion sends. Both label
// dimensions that have leaked are request-controlled — the path and
// the method — so both vary here and one driver sends them.
type probe struct {
	method string
	path   string
}

// drive sends every probe and returns how many responses carried each
// status code.
func drive(t *testing.T, h http.Handler, probes []probe) map[int]int {
	t.Helper()

	codes := make(map[int]int)

	for _, p := range probes {
		req := httptest.NewRequestWithContext(
			t.Context(), p.method, p.path, nil,
		)
		w := httptest.NewRecorder()
		h.ServeHTTP(w, req)
		codes[w.Code]++
	}

	return codes
}

// drivePaths sends one POST per supplied path.
func drivePaths(
	t *testing.T,
	h http.Handler,
	paths []string,
) map[int]int {
	t.Helper()

	probes := make([]probe, 0, len(paths))

	for _, p := range paths {
		probes = append(
			probes, probe{method: http.MethodPost, path: p},
		)
	}

	return drive(t, h, probes)
}

// receiverPaths returns n distinct /webhook/ paths, each naming a
// fresh UUID exactly as an unauthenticated flood would.
func receiverPaths(n int) []string {
	paths := make([]string, 0, n)

	for range n {
		paths = append(paths, "/webhook/"+uuid.NewString())
	}

	return paths
}

// gatherMetrics returns the registry's current families, failing the
// test if gathering does.
func gatherMetrics(
	t *testing.T,
	reg *prometheus.Registry,
) []*dto.MetricFamily {
	t.Helper()

	families, err := reg.Gather()
	require.NoError(t, err)

	return families
}

// labelValue returns the named label from a gathered metric.
func labelValue(m *dto.Metric, name string) string {
	for _, pair := range m.GetLabel() {
		if pair.GetName() == name {
			return pair.GetValue()
		}
	}

	return ""
}

// handlerLabels returns the set of distinct `handler` label values
// across every gathered series.
func handlerLabels(families []*dto.MetricFamily) map[string]struct{} {
	seen := make(map[string]struct{})

	for _, fam := range families {
		for _, m := range fam.GetMetric() {
			seen[labelValue(m, "handler")] = struct{}{}
		}
	}

	return seen
}

// seriesCount is the number of distinct label sets held across every
// family: the quantity that grew without bound and was never
// reclaimed.
func seriesCount(families []*dto.MetricFamily) int {
	total := 0

	for _, fam := range families {
		total += len(fam.GetMetric())
	}

	return total
}

// keys returns the members of a set, for assertion messages.
func keys(set map[string]struct{}) []string {
	out := make([]string, 0, len(set))

	for k := range set {
		out = append(out, k)
	}

	return out
}

// TestMetrics_DistinctReceiverPathsMintOneLabelSet is the direct
// assertion the issue asks for: N requests to N distinct
// /webhook/<uuid> paths must produce exactly ONE handler label, the
// route pattern. Before the fix this produced N of them.
func TestMetrics_DistinctReceiverPathsMintOneLabelSet(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, generousReceiverLimit)

	paths := receiverPaths(metricsProbePaths)
	codes := drivePaths(t, h, paths)
	require.Equal(
		t, metricsProbePaths, codes[http.StatusNotFound],
		"every invented UUID should have reached the receiver",
	)

	families := gatherMetrics(t, reg)
	labels := handlerLabels(families)

	assert.Equal(
		t,
		map[string]struct{}{
			receiverRoutePattern:            {},
			middleware.InflightHandlerConst: {},
		},
		labels,
		"receiver traffic must collapse onto the route pattern",
	)

	// The scrape must not republish the UUIDs it was driven with.
	// They are the receiver's only credential.
	for _, p := range paths {
		id := strings.TrimPrefix(p, "/webhook/")
		for label := range labels {
			assert.NotContains(
				t, label, id,
				"an entrypoint UUID reached a metrics label",
			)
		}
	}
}

// TestMetrics_SeriesCountIsFlatUnderAFlood pins the property the
// issue measured against a live instance: driving thousands more
// distinct paths must not add series. The first batch establishes
// every label set the route can produce; the second must land on
// exactly those.
func TestMetrics_SeriesCountIsFlatUnderAFlood(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, generousReceiverLimit)

	drivePaths(t, h, receiverPaths(metricsProbePaths))
	before := seriesCount(gatherMetrics(t, reg))

	drivePaths(t, h, receiverPaths(metricsProbePaths*4))
	after := seriesCount(gatherMetrics(t, reg))

	assert.Equal(
		t, before, after,
		"a flood of distinct paths must not mint series",
	)
}

// TestMetrics_RateLimitedRequestsCarryTheRoutePattern covers the
// majority case: most of the leaked series were 429s. Those requests
// never reach a handler, so they take a different path through the
// stack — but chi has already matched the route by the time the
// route-level limiter rejects them, and the recording happens after
// the whole chain returns, so they must land on the pattern too.
func TestMetrics_RateLimitedRequestsCarryTheRoutePattern(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, tightReceiverLimit)

	codes := drivePaths(t, h, receiverPaths(metricsProbePaths))
	require.Positive(
		t, codes[http.StatusTooManyRequests],
		"the probe must actually exhaust the aggregate limiter",
	)

	families := gatherMetrics(t, reg)

	assert.Equal(
		t,
		map[string]struct{}{
			receiverRoutePattern:            {},
			middleware.InflightHandlerConst: {},
		},
		handlerLabels(families),
		"rejected requests must collapse onto the route pattern",
	)

	rejected := 0

	for _, fam := range families {
		for _, m := range fam.GetMetric() {
			if labelValue(m, "code") != "429" {
				continue
			}

			rejected++

			assert.Equal(
				t, receiverRoutePattern,
				labelValue(m, "handler"),
				"a 429 series carried a non-pattern handler",
			)
		}
	}

	assert.Positive(
		t, rejected, "no 429 series was recorded at all",
	)
}

// TestMetrics_UnmatchedPathsCollapseToTheSentinel decides and pins the
// unmatched-route case. A path matching no route has no pattern, so
// it carries the same fixed sentinel the access log uses. Without
// that, an unmatched flood leaks exactly as the receiver did.
func TestMetrics_UnmatchedPathsCollapseToTheSentinel(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, generousReceiverLimit)

	paths := make([]string, 0, metricsProbePaths)

	for i := range metricsProbePaths {
		id := uuid.NewString()

		// Two shapes: one matching no prefix at all, and one under
		// the receiver prefix but with a segment count the pattern
		// cannot match.
		if i%2 == 0 {
			paths = append(paths, "/"+id)
		} else {
			paths = append(paths, "/webhook/"+id+"/"+id)
		}
	}

	codes := drivePaths(t, h, paths)
	require.Equal(
		t, metricsProbePaths, codes[http.StatusNotFound],
		"every probe path should have gone unmatched",
	)

	labels := handlerLabels(gatherMetrics(t, reg))

	assert.Equal(
		t,
		map[string]struct{}{
			middleware.UnmatchedRouteConst:  {},
			middleware.InflightHandlerConst: {},
		},
		labels,
		"unmatched paths must collapse onto one sentinel, got %v",
		keys(labels),
	)
}

// TestMetrics_InflightGaugeIsAggregateAndBalanced pins the one metric
// that cannot carry a pattern. It is incremented before routing and
// decremented after, so it gets a fixed label -- and the two calls
// must therefore agree, leaving the gauge at zero once the traffic
// has drained rather than stuck above it.
func TestMetrics_InflightGaugeIsAggregateAndBalanced(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, generousReceiverLimit)

	drivePaths(t, h, receiverPaths(metricsProbePaths))

	var inflight []*dto.Metric

	for _, fam := range gatherMetrics(t, reg) {
		if strings.HasSuffix(fam.GetName(), "requests_inflight") {
			inflight = fam.GetMetric()
		}
	}

	require.Len(
		t, inflight, 1,
		"the inflight gauge must hold exactly one series",
	)
	assert.Equal(
		t, middleware.InflightHandlerConst,
		labelValue(inflight[0], "handler"),
	)
	assert.InDelta(
		t, 0.0, inflight[0].GetGauge().GetValue(), 0.0,
		"the gauge must balance back to zero",
	)
}

// TestMetrics_StatusAndSizeStillRecorded guards the response-writer
// interceptor the recording middleware wraps around every request.
// The handler label changed; what the interceptor reports must not
// have.
func TestMetrics_StatusAndSizeStillRecorded(t *testing.T) {
	t.Parallel()

	h, reg := metricsTestRouter(t, generousReceiverLimit)

	req := httptest.NewRequestWithContext(
		t.Context(), http.MethodGet, okRoute, nil,
	)
	w := httptest.NewRecorder()
	h.ServeHTTP(w, req)

	require.Equal(t, http.StatusOK, w.Code)
	require.Equal(t, okBody, w.Body.String())

	var size *dto.Metric

	for _, fam := range gatherMetrics(t, reg) {
		if !strings.HasSuffix(fam.GetName(), "response_size_bytes") {
			continue
		}

		for _, m := range fam.GetMetric() {
			if labelValue(m, "handler") == okRoute {
				size = m
			}
		}
	}

	require.NotNil(
		t, size, "no response size series for the static route",
	)
	assert.Equal(t, "200", labelValue(size, "code"))
	assert.Equal(t, uint64(1), size.GetHistogram().GetSampleCount())
	assert.InDelta(
		t, float64(len(okBody)),
		size.GetHistogram().GetSampleSum(), 0.0,
		"the interceptor must still count written bytes",
	)
}