Merge pull request #9 from xray33/feature/ability_to_skip_http_metric…

…s_reporting

Ability to skip HTTP metrics reporting

httpserver/middleware/context.go

@@ -22,6 +22,7 @@ const (
 	ctxKeyLoggingParams
 	ctxKeyTraceID
 	ctxKeyRequestStartTime
+	ctxKeyHTTPMetricsStatus
 )
 
 func getStringFromContext(ctx context.Context, key ctxKey) string {
@@ -100,3 +101,27 @@ func GetRequestStartTimeFromContext(ctx context.Context) time.Time {
 	startTime, _ := ctx.Value(ctxKeyRequestStartTime).(time.Time)
 	return startTime
 }
+
+// NewContextWithHTTPMetricsEnabled creates a new context with special flag which can toggle HTTP metrics status.
+func NewContextWithHTTPMetricsEnabled(ctx context.Context) context.Context {
+	status := true
+	return context.WithValue(ctx, ctxKeyHTTPMetricsStatus, &status)
+}
+
+// IsHTTPMetricsEnabledInContext checks whether HTTP metrics are enabled in the context.
+func IsHTTPMetricsEnabledInContext(ctx context.Context) bool {
+	value := ctx.Value(ctxKeyHTTPMetricsStatus)
+	if value == nil {
+		return false
+	}
+	return *value.(*bool)
+}
+
+// DisableHTTPMetricsInContext disables HTTP metrics processing in the context.
+func DisableHTTPMetricsInContext(ctx context.Context) {
+	value := ctx.Value(ctxKeyHTTPMetricsStatus)
+	if value == nil {
+		return
+	}
+	*value.(*bool) = false
+}

httpserver/middleware/metrics.go

@@ -214,8 +214,14 @@ func (h *httpRequestMetricsHandler) ServeHTTP(rw http.ResponseWriter, r *http.Re
 	inFlightGauge.Inc()
 	defer inFlightGauge.Dec()
 
+	r = r.WithContext(NewContextWithHTTPMetricsEnabled(r.Context()))
+
 	wrw := WrapResponseWriterIfNeeded(rw, r.ProtoMajor)
 	defer func() {
+		if !IsHTTPMetricsEnabledInContext(r.Context()) {
+			return
+		}
+
 		if reqInfo.routePattern == "" {
 			reqInfo.routePattern = h.getRoutePattern(r)
 		}

httpserver/middleware/metrics_test.go

@@ -28,6 +28,13 @@ func (h *mockHTTPRequestMetricsNextHandler) ServeHTTP(rw http.ResponseWriter, r
 	rw.WriteHeader(h.statusCodeToReturn)
 }
 
+type mockHTTPRequestMetricsDisabledHandler struct{}
+
+func (h *mockHTTPRequestMetricsDisabledHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	DisableHTTPMetricsInContext(r.Context())
+	rw.WriteHeader(http.StatusOK)
+}
+
 func TestHttpRequestMetricsHandler_ServeHTTP(t *testing.T) {
 	makeLabels := func(method, routePattern, uaType, statusCode string) prometheus.Labels {
 		return prometheus.Labels{
@@ -173,4 +180,18 @@ func TestHttpRequestMetricsHandler_ServeHTTP(t *testing.T) {
 			testutil.AssertSamplesCountInHistogram(t, hist, 1)
 		}
 	})
+
+	t.Run("not collect if disabled", func(t *testing.T) {
+		collector := NewHTTPRequestMetricsCollector()
+		next := &mockHTTPRequestMetricsDisabledHandler{}
+		req := httptest.NewRequest(http.MethodGet, "/hello", nil)
+		req.Header.Set("User-Agent", "http-client")
+		resp := httptest.NewRecorder()
+		h := HTTPRequestMetrics(collector, getRoutePattern)(next)
+		h.ServeHTTP(resp, req)
+		assert.Equal(t, http.StatusOK, resp.Code)
+		labels := makeLabels(http.MethodGet, "/hello", "http-client", "200")
+		hist := collector.Durations.With(labels).(prometheus.Histogram)
+		testutil.AssertSamplesCountInHistogram(t, hist, 0)
+	})
 }