grafana/pkg/middleware/request_tracing.go

package middleware

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"

	"github.com/grafana/grafana/pkg/web"
)

type contextKey struct{}

var routeOperationNameKey = contextKey{}

// ProvideRouteOperationName creates a named middleware responsible for populating
// the context with the route operation name that can be used later in the request pipeline.
// 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)
	}
}

// RouteOperationNameFromContext receives the route operation name from context, if set.
func RouteOperationNameFromContext(ctx context.Context) (string, bool) {
	if val := ctx.Value(routeOperationNameKey); val != nil {
		op, ok := val.(string)
		return op, ok
	}

	return "", false
}

func RequestTracing() web.Handler {
	return func(res http.ResponseWriter, req *http.Request, c *web.Context) {
		if strings.HasPrefix(c.Req.URL.Path, "/public/") ||
			c.Req.URL.Path == "robots.txt" {
			c.Next()
			return
		}

		rw := res.(web.ResponseWriter)

		tracer := opentracing.GlobalTracer()
		wireContext, _ := tracer.Extract(opentracing.HTTPHeaders, opentracing.HTTPHeadersCarrier(req.Header))
		span := tracer.StartSpan(fmt.Sprintf("HTTP %s %s", req.Method, req.URL.Path), ext.RPCServerOption(wireContext))

		ctx := opentracing.ContextWithSpan(req.Context(), span)
		c.Req = req.WithContext(ctx)
		c.Map(c.Req)

		c.Next()

		// Only call span.Finish when a route operation name have been set,
		// meaning that not set the span would not be reported.
		if routeOperation, exists := RouteOperationNameFromContext(c.Req.Context()); exists {
			defer span.Finish()
			span.SetOperationName(fmt.Sprintf("HTTP %s %s", req.Method, routeOperation))
		}

		status := rw.Status()

		ext.HTTPStatusCode.Set(span, uint16(status))
		ext.HTTPUrl.Set(span, req.RequestURI)
		ext.HTTPMethod.Set(span, req.Method)
		if status >= 400 {
			ext.Error.Set(span, true)
		}
	}
}