From ab59f3cb1a71e3f7b31fe71b0d19b5754733c640 Mon Sep 17 00:00:00 2001
From: Carl Bergquist <carl.bergquist@gmail.com>
Date: Thu, 25 Aug 2022 10:11:27 +0200
Subject: [PATCH] metrics: add operation name used in metrics/tracing for
 middlewares (#53949)

Signed-off-by: bergquist <carl.bergquist@gmail.com>
---
 pkg/api/http_server.go            |  2 +-
 pkg/middleware/request_test.go    | 18 ++++++++++++++++++
 pkg/middleware/request_tracing.go | 13 +++++++++++--
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/pkg/api/http_server.go b/pkg/api/http_server.go
index 7e18417e491..7887086efc9 100644
--- a/pkg/api/http_server.go
+++ b/pkg/api/http_server.go
@@ -525,7 +525,7 @@ func (hs *HTTPServer) applyRoutes() {
 	// then add view routes & api routes
 	hs.RouteRegister.Register(hs.web, hs.namedMiddlewares...)
 	// lastly not found route
-	hs.web.NotFound(middleware.ReqSignedIn, hs.NotFoundHandler)
+	hs.web.NotFound(middleware.ProvideRouteOperationName("notfound"), middleware.ReqSignedIn, hs.NotFoundHandler)
 }
 
 func (hs *HTTPServer) addMiddlewaresAndStaticRoutes() {
diff --git a/pkg/middleware/request_test.go b/pkg/middleware/request_test.go
index 16b41b9f231..13303025a9b 100644
--- a/pkg/middleware/request_test.go
+++ b/pkg/middleware/request_test.go
@@ -38,3 +38,21 @@ func TestCanGetRouteNameFromContext(t *testing.T) {
 		assert.Equal(t, tc.expected, handler)
 	}
 }
+
+func TestOperationNameCanOnlyBeSetOnce(t *testing.T) {
+	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://grafana.com", nil)
+
+	// set the the initial operation name
+	req = addRouteNameToContext(req, "first")
+
+	// check that the operation name is set correctly
+	value, exists := routeOperationName(req)
+	assert.True(t, exists, "route name should exist")
+	assert.Equal(t, "first", value)
+
+	// check that it cannot be overwritten
+	req = addRouteNameToContext(req, "second")
+	value, exists = routeOperationName(req)
+	assert.True(t, exists, "route name should exist")
+	assert.Equal(t, "first", value)
+}
diff --git a/pkg/middleware/request_tracing.go b/pkg/middleware/request_tracing.go
index 64fe6311eb6..e86cc4f3388 100644
--- a/pkg/middleware/request_tracing.go
+++ b/pkg/middleware/request_tracing.go
@@ -27,11 +27,20 @@ var routeOperationNameKey = contextKey{}
 // Implements routing.RegisterNamedMiddleware.
 func ProvideRouteOperationName(name string) web.Handler {
 	return func(res http.ResponseWriter, req *http.Request, c *web.Context) {
-		ctx := context.WithValue(c.Req.Context(), routeOperationNameKey, name)
-		c.Req = c.Req.WithContext(ctx)
+		c.Req = addRouteNameToContext(c.Req, name)
 	}
 }
 
+func addRouteNameToContext(req *http.Request, operationName string) *http.Request {
+	// don't set route name if it's set
+	if _, exists := routeOperationName(req); exists {
+		return req
+	}
+
+	ctx := context.WithValue(req.Context(), routeOperationNameKey, operationName)
+	return req.WithContext(ctx)
+}
+
 var unnamedHandlers = []struct {
 	pathPattern *regexp.Regexp
 	handler     string