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

Ref #68480

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

pkg/util/proxyutil/reverse_proxy_test.go

@@ -3,6 +3,7 @@ package proxyutil
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -11,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/middleware/requestmeta"
 	"github.com/grafana/grafana/pkg/services/contexthandler"
 	"github.com/grafana/grafana/pkg/setting"
 )
@@ -97,7 +99,7 @@ func TestReverseProxy(t *testing.T) {
 		require.NoError(t, resp.Body.Close())
 	})
 
-	t.Run("Error handling should convert status codes depending on what kind of error it is", func(t *testing.T) {
+	t.Run("Error handling should convert status codes depending on what kind of error it is and set downstream status source", func(t *testing.T) {
 		timedOutTransport := http.DefaultTransport.(*http.Transport)
 		timedOutTransport.ResponseHeaderTimeout = time.Millisecond
 
@@ -136,7 +138,12 @@ func TestReverseProxy(t *testing.T) {
 				}))
 				t.Cleanup(upstream.Close)
 				rec := httptest.NewRecorder()
+
+				ctx := requestmeta.SetRequestMetaData(context.Background(), requestmeta.RequestMetaData{
+					StatusSource: requestmeta.StatusSourceServer,
+				})
 				req := httptest.NewRequest(http.MethodGet, upstream.URL, nil)
+				req = req.WithContext(ctx)
 
 				rp := NewReverseProxy(
 					log.New("test"),
@@ -151,6 +158,55 @@ func TestReverseProxy(t *testing.T) {
 				resp := rec.Result()
 				require.Equal(t, tc.expectedStatusCode, resp.StatusCode)
 				require.NoError(t, resp.Body.Close())
+
+				rmd := requestmeta.GetRequestMetaData(ctx)
+				require.Equal(t, requestmeta.StatusSourceDownstream, rmd.StatusSource)
+			})
+		}
+	})
+
+	t.Run("5xx response status codes should set downstream status source", func(t *testing.T) {
+		testCases := []struct {
+			status         int
+			expectedSource requestmeta.StatusSource
+		}{
+			{status: http.StatusOK, expectedSource: requestmeta.StatusSourceServer},
+			{status: http.StatusBadRequest, expectedSource: requestmeta.StatusSourceServer},
+			{status: http.StatusForbidden, expectedSource: requestmeta.StatusSourceServer},
+			{status: http.StatusUnauthorized, expectedSource: requestmeta.StatusSourceServer},
+			{status: http.StatusInternalServerError, expectedSource: requestmeta.StatusSourceDownstream},
+			{status: http.StatusBadGateway, expectedSource: requestmeta.StatusSourceDownstream},
+			{status: http.StatusGatewayTimeout, expectedSource: requestmeta.StatusSourceDownstream},
+			{status: 599, expectedSource: requestmeta.StatusSourceDownstream},
+		}
+
+		for _, tc := range testCases {
+			t.Run(fmt.Sprintf("status %d => source %s ", tc.status, tc.expectedSource), func(t *testing.T) {
+				upstream := newUpstreamServer(t, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+					w.WriteHeader(tc.status)
+				}))
+				t.Cleanup(upstream.Close)
+				rec := httptest.NewRecorder()
+
+				ctx := requestmeta.SetRequestMetaData(context.Background(), requestmeta.RequestMetaData{
+					StatusSource: requestmeta.StatusSourceServer,
+				})
+				req := httptest.NewRequest(http.MethodGet, upstream.URL, nil)
+				req = req.WithContext(ctx)
+
+				rp := NewReverseProxy(
+					log.New("test"),
+					func(req *http.Request) {},
+				)
+				require.NotNil(t, rp)
+				rp.ServeHTTP(rec, req)
+
+				resp := rec.Result()
+				require.Equal(t, tc.status, resp.StatusCode)
+				require.NoError(t, resp.Body.Close())
+
+				rmd := requestmeta.GetRequestMetaData(ctx)
+				require.Equal(t, tc.expectedSource, rmd.StatusSource)
 			})
 		}
 	})