Instrumentation: Add status_source label to request metrics/logs (#74114)

Ref #68480

Co-authored-by: Giuseppe Guerra <giuseppe.guerra@grafana.com>

pkg/middleware/loggermw/logger.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/grafana/grafana/pkg/middleware"
+	"github.com/grafana/grafana/pkg/middleware/requestmeta"
 
 	"github.com/grafana/grafana/pkg/util/errutil"
 
@@ -127,6 +128,11 @@ func (l *loggerImpl) prepareLogParams(c *contextmodel.ReqContext, duration time.
 		logParams = append(logParams, "handler", handler)
 	}
 
+	if l.flags.IsEnabled(featuremgmt.FlagRequestInstrumentationStatusSource) {
+		rmd := requestmeta.GetRequestMetaData(c.Req.Context())
+		logParams = append(logParams, "status_source", rmd.StatusSource)
+	}
+
 	logParams = append(logParams, errorLogParams(c.Error)...)
 
 	return logParams, lvl

pkg/middleware/request_metrics.go

@@ -36,6 +36,11 @@ func RequestMetrics(features featuremgmt.FeatureToggles, cfg *setting.Cfg, promR
 	)
 
 	histogramLabels := []string{"handler", "status_code", "method"}
+
+	if features.IsEnabled(featuremgmt.FlagRequestInstrumentationStatusSource) {
+		histogramLabels = append(histogramLabels, "status_source")
+	}
+
 	if cfg.MetricsIncludeTeamLabel {
 		histogramLabels = append(histogramLabels, "team")
 	}
@@ -80,8 +85,13 @@ func RequestMetrics(features featuremgmt.FeatureToggles, cfg *setting.Cfg, promR
 			}
 
 			labelValues := []string{handler, code, r.Method}
+			rmd := requestmeta.GetRequestMetaData(r.Context())
+
+			if features.IsEnabled(featuremgmt.FlagRequestInstrumentationStatusSource) {
+				labelValues = append(labelValues, string(rmd.StatusSource))
+			}
+
 			if cfg.MetricsIncludeTeamLabel {
-				rmd := requestmeta.GetRequestMetaData(r.Context())
 				labelValues = append(labelValues, rmd.Team)
 			}
 

pkg/middleware/requestmeta/request_metadata.go

@@ -13,10 +13,18 @@ const (
 	TeamCore     = "core"
 )
 
+type StatusSource string
+
+const (
+	StatusSourceServer     StatusSource = "server"
+	StatusSourceDownstream StatusSource = "downstream"
+)
+
 type rMDContextKey struct{}
 
 type RequestMetaData struct {
-	Team string
+	Team         string
+	StatusSource StatusSource
 }
 
 var requestMetaDataContextKey = rMDContextKey{}
@@ -27,8 +35,7 @@ func SetupRequestMetadata() web.Middleware {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			rmd := defaultRequestMetadata()
-
-			ctx := context.WithValue(r.Context(), requestMetaDataContextKey, rmd)
+			ctx := SetRequestMetaData(r.Context(), rmd)
 			*r = *r.WithContext(ctx)
 
 			next.ServeHTTP(w, r)
@@ -46,7 +53,13 @@ func GetRequestMetaData(ctx context.Context) *RequestMetaData {
 		return value
 	}
 
-	return defaultRequestMetadata()
+	rmd := defaultRequestMetadata()
+	return &rmd
+}
+
+// SetRequestMetaData sets the request metadata for the context.
+func SetRequestMetaData(ctx context.Context, rmd RequestMetaData) context.Context {
+	return context.WithValue(ctx, requestMetaDataContextKey, &rmd)
 }
 
 // SetOwner returns an `web.Handler` that sets the team name for an request.
@@ -57,8 +70,31 @@ func SetOwner(team string) web.Handler {
 	}
 }
 
-func defaultRequestMetadata() *RequestMetaData {
-	return &RequestMetaData{
-		Team: TeamCore,
+// WithDownstreamStatusSource sets the StatusSource field of the [RequestMetaData] for the
+// context to [StatusSourceDownstream].
+func WithDownstreamStatusSource(ctx context.Context) {
+	v := GetRequestMetaData(ctx)
+	v.StatusSource = StatusSourceDownstream
+}
+
+// WithStatusSource sets the StatusSource field of the [RequestMetaData] for the
+// context based on the provided statusCode.
+// If statusCode >= 500 then [StatusSourceDownstream].
+// If statusCode < 500 then [StatusSourceServer].
+func WithStatusSource(ctx context.Context, statusCode int) {
+	v := GetRequestMetaData(ctx)
+
+	if statusCode >= 500 {
+		v.StatusSource = StatusSourceDownstream
+		return
+	}
+
+	v.StatusSource = StatusSourceServer
+}
+
+func defaultRequestMetadata() RequestMetaData {
+	return RequestMetaData{
+		Team:         TeamCore,
+		StatusSource: StatusSourceServer,
 	}
 }

pkg/middleware/requestmeta/request_metadata_test.go

@@ -0,0 +1,47 @@
+package requestmeta
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestStatusSource(t *testing.T) {
+	ctx := context.Background()
+	ctx = SetRequestMetaData(ctx, defaultRequestMetadata())
+	rmd := GetRequestMetaData(ctx)
+	require.Equal(t, StatusSourceServer, rmd.StatusSource)
+
+	WithDownstreamStatusSource(ctx)
+	rmd = GetRequestMetaData(ctx)
+	require.Equal(t, StatusSourceDownstream, rmd.StatusSource)
+}
+
+func TestWithStatusSource(t *testing.T) {
+	tcs := []struct {
+		status         int
+		expectedSource StatusSource
+	}{
+		{status: http.StatusOK, expectedSource: StatusSourceServer},
+		{status: http.StatusBadRequest, expectedSource: StatusSourceServer},
+		{status: http.StatusForbidden, expectedSource: StatusSourceServer},
+		{status: http.StatusUnauthorized, expectedSource: StatusSourceServer},
+		{status: http.StatusInternalServerError, expectedSource: StatusSourceDownstream},
+		{status: http.StatusBadGateway, expectedSource: StatusSourceDownstream},
+		{status: http.StatusGatewayTimeout, expectedSource: StatusSourceDownstream},
+		{status: 599, expectedSource: StatusSourceDownstream},
+	}
+
+	for _, tc := range tcs {
+		t.Run(fmt.Sprintf("status %d => source %s ", tc.status, tc.expectedSource), func(t *testing.T) {
+			ctx := context.Background()
+			ctx = SetRequestMetaData(ctx, defaultRequestMetadata())
+			WithStatusSource(ctx, tc.status)
+			rmd := GetRequestMetaData(ctx)
+			require.Equal(t, tc.expectedSource, rmd.StatusSource)
+		})
+	}
+}