Wrap the defer in a function to make it defer after the return rather than after the if block. Add a unit test to validate we're tracking time correctly.

Signed-off-by: Tyler Reid <[email protected]>

Author: Tyler Reid

pkg/ruler/compat.go

@@ -151,7 +151,7 @@ func MetricsQueryFunc(qf rules.QueryFunc, queries, failedQueries prometheus.Coun
 		// If we've been passed a counter vec we want to record the wall time spent executing this request.
 		if queryTime != nil {
 			startTime = time.Now()
-			defer queryTime.WithLabelValues(userID).Add(float64(time.Since(startTime)))
+			defer func() { queryTime.WithLabelValues(userID).Add(float64(time.Since(startTime))) }()
 		}
 
 		result, err := qf(ctx, qs, t)

pkg/ruler/compat_test.go

@@ -231,8 +231,7 @@ func TestMetricsQueryFuncErrors(t *testing.T) {
 			mockFunc := func(ctx context.Context, q string, t time.Time) (promql.Vector, error) {
 				return promql.Vector{}, tc.returnedError
 			}
-
-			qf := MetricsQueryFunc(mockFunc, queries, failures, queryTime, "user")
+			qf := MetricsQueryFunc(mockFunc, queries, failures, queryTime, "userID")
 
 			_, err := qf(context.Background(), "test", time.Now())
 			require.Equal(t, tc.returnedError, err)
@@ -242,3 +241,21 @@ func TestMetricsQueryFuncErrors(t *testing.T) {
 		})
 	}
 }
+
+func TestMetricsQueryFuncMetrics(t *testing.T) {
+	queries := prometheus.NewCounter(prometheus.CounterOpts{})
+	failures := prometheus.NewCounter(prometheus.CounterOpts{})
+	queryTime := prometheus.NewCounterVec(prometheus.CounterOpts{}, []string{"user"})
+
+	mockFunc := func(ctx context.Context, q string, t time.Time) (promql.Vector, error) {
+		time.Sleep(1 * time.Millisecond)
+		return promql.Vector{}, nil
+	}
+	qf := MetricsQueryFunc(mockFunc, queries, failures, queryTime, "userID")
+
+	_, _ = qf(context.Background(), "test", time.Now())
+
+	require.Equal(t, 1, int(testutil.ToFloat64(queries)))
+	require.Equal(t, 0, int(testutil.ToFloat64(failures)))
+	require.LessOrEqual(t, float64(1*time.Millisecond), testutil.ToFloat64(queryTime.WithLabelValues("userID")))
+}