feat: enrich trace span information (#76)

README.md

@@ -127,8 +127,10 @@ Shield provides a fully-featured GRPC and HTTP API to interact with Shield serve
   <summary>Dependencies:</summary>
 
     - Git
-    - Go 1.17 or above
-    - PostgreSQL 13.2 or above
+    - Go 1.21 or above
+    - PostgreSQL 15 or above
+      - pg_partman 4.7.4
+      - pg_cron 1.6
 
 </details>
 
@@ -159,6 +161,11 @@ cp internal/server/config.yaml config.yaml
 
 Run database migrations
 
+Add these config in postgresql.conf
+```
+cron.database_name = {database name}
+shared_preload_libraries = 'pg_cron'
+```
 ```
 ./shield server migrate -c config.yaml
 ```

cmd/serve_proxy.go

@@ -21,12 +21,10 @@ import (
 	"github.com/goto/shield/internal/proxy/middleware/authz"
 	"github.com/goto/shield/internal/proxy/middleware/basic_auth"
 	"github.com/goto/shield/internal/proxy/middleware/observability"
+	"github.com/goto/shield/internal/proxy/middleware/otelpostprocessor"
 	"github.com/goto/shield/internal/proxy/middleware/prefix"
 	"github.com/goto/shield/internal/proxy/middleware/rulematch"
 	"github.com/goto/shield/internal/store/blob"
-	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
-	semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
-	"go.opentelemetry.io/otel/trace"
 )
 
 func serveProxies(
@@ -108,19 +106,8 @@ func buildMiddlewarePipeline(
 	casbinAuthz := authz.New(logger, prefixWare, userIDHeaderKey, resourceService, userService, groupService)
 	basicAuthn := basic_auth.New(logger, casbinAuthz)
 	attributeExtractor := attributes.New(logger, basicAuthn, identityProxyHeaderKey, projectService)
-	otelMiddleware := otelhttp.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		route := r.URL.Path
-		attr := semconv.HTTPRouteKey.String(route)
-
-		span := trace.SpanFromContext(r.Context())
-		span.SetAttributes(attr)
-
-		labeler, _ := otelhttp.LabelerFromContext(r.Context())
-		labeler.Add(attr)
-
-		attributeExtractor.ServeHTTP(w, r)
-	}), "")
-	matchWare := rulematch.New(logger, otelMiddleware, rulematch.NewRouteMatcher(ruleService))
+	otelPostProcessor := otelpostprocessor.New(attributeExtractor)
+	matchWare := rulematch.New(logger, otelPostProcessor, rulematch.NewRouteMatcher(ruleService))
 	observability := observability.New(logger, matchWare)
 	return observability
 }

internal/proxy/middleware/context.go

@@ -47,8 +47,8 @@ func ExtractMiddleware(r *http.Request, name string) (rule.MiddlewareSpec, bool)
 	return rl.Middlewares.Get(name)
 }
 
-func EnrichRequestWithMuxRouteAndVars(r *http.Request, route *mux.Route, vars map[string]string) {
-	*r = *r.WithContext(httputil.SetContextWithMuxRouteAndVars(r.Context(), route, vars))
+func EnrichRequestWithMuxRoute(r *http.Request, route *mux.Route) {
+	*r = *r.WithContext(httputil.SetContextWithMuxRoute(r.Context(), route))
 }
 
 func EnrichPathParams(r *http.Request, params map[string]string) {

internal/proxy/middleware/observability/observability.go

@@ -5,6 +5,9 @@ import (
 	"strings"
 
 	"github.com/goto/shield/internal/proxy/middleware"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	"go.opentelemetry.io/otel/trace"
 
 	"github.com/goto/salt/log"
 	"github.com/rs/xid"
@@ -16,14 +19,24 @@ const (
 )
 
 type Ware struct {
-	log  *log.Zap
-	next http.Handler
+	log         *log.Zap
+	otelHandler http.Handler
 }
 
 func New(log *log.Zap, next http.Handler) *Ware {
 	return &Ware{
-		log:  log,
-		next: next,
+		log: log,
+		otelHandler: otelhttp.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attr := semconv.HTTPRouteKey.String(r.URL.Path)
+
+			span := trace.SpanFromContext(r.Context())
+			span.SetAttributes(attr)
+
+			labeler, _ := otelhttp.LabelerFromContext(r.Context())
+			labeler.Add(attr)
+
+			next.ServeHTTP(w, r)
+		}), ""),
 	}
 }
 
@@ -46,7 +59,8 @@ func (m *Ware) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		zap.String("request_id", reqID),
 	)
 	req = req.WithContext(ctx)
-	m.next.ServeHTTP(rw, req)
+
+	m.otelHandler.ServeHTTP(rw, req)
 }
 
 func setRequestID(req *http.Request) string {

internal/proxy/middleware/otelpostprocessor/otelpostprocessor.go

@@ -0,0 +1,46 @@
+package otelpostprocessor
+
+import (
+	"net/http"
+
+	"github.com/goto/shield/internal/proxy/middleware"
+	"github.com/goto/shield/pkg/httputil"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type Ware struct {
+	next http.Handler
+}
+
+func New(next http.Handler) *Ware {
+	return &Ware{
+		next: next,
+	}
+}
+
+func (m Ware) Info() *middleware.MiddlewareInfo {
+	return &middleware.MiddlewareInfo{
+		Name:        "_otel_postprocessor",
+		Description: "post process otel metrics and traces",
+	}
+}
+
+func (m *Ware) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	routeName := httputil.MuxRouteName(r)
+	labeler, _ := otelhttp.LabelerFromContext(r.Context())
+	labeler.Add(semconv.HTTPRouteKey.String(routeName))
+
+	span := trace.SpanFromContext(r.Context())
+	span.SetName(m.SpanName(r, routeName))
+
+	ctx := trace.ContextWithSpan(r.Context(), span)
+	r = r.WithContext(ctx)
+
+	m.next.ServeHTTP(rw, r)
+}
+
+func (w *Ware) SpanName(r *http.Request, routeName string) string {
+	return r.Method + " " + routeName
+}

internal/proxy/middleware/otelpostprocessor/otelpostprocessor_test.go

@@ -0,0 +1,31 @@
+package otelpostprocessor_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/goto/shield/internal/proxy/middleware"
+	"github.com/goto/shield/internal/proxy/middleware/otelpostprocessor"
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+)
+
+func TestWare_ServeHTTP(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		labels, exist := otelhttp.LabelerFromContext(r.Context())
+		assert.NotNil(t, labels)
+		assert.True(t, exist)
+	})
+
+	op := otelpostprocessor.New(next)
+	o := otelhttp.NewHandler(op, "")
+
+	assert.Equal(t, &middleware.MiddlewareInfo{
+		Name:        "_otel_postprocessor",
+		Description: "post process otel metrics and traces",
+	}, op.Info())
+
+	req := httptest.NewRequest("GET", "http://testing", nil)
+	o.ServeHTTP(httptest.NewRecorder(), req)
+}

internal/proxy/middleware/rulematch/route_matcher.go

@@ -26,7 +26,7 @@ func (r RouteMatcher) Match(req *http.Request) (*rule.Rule, error) {
 			route := router.NewRoute().Path(rule.Frontend.URL).Methods(rule.Frontend.Method)
 			routeMatcher := mux.RouteMatch{}
 			if route.Match(req, &routeMatcher) {
-				middleware.EnrichRequestWithMuxRouteAndVars(req, routeMatcher.Route, routeMatcher.Vars)
+				middleware.EnrichRequestWithMuxRoute(req, routeMatcher.Route)
 				middleware.EnrichPathParams(req, routeMatcher.Vars)
 				return &rule, nil
 			}

internal/proxy/roundtripper.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/goto/shield/internal/proxy/hook"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 
 	"github.com/goto/salt/log"
 	"go.uber.org/zap"
@@ -22,8 +23,8 @@ type h2cTransportWrapper struct {
 	// proxy for http & h2
 	// Reference: https://sourcegraph.com/github.com/tsenart/vegeta/-/blob/lib/attack.go?L206:6#tab=references
 
-	httpTransport *http.Transport
-	grpcTransport *http2.Transport
+	httpTransport *otelhttp.Transport
+	grpcTransport *otelhttp.Transport
 
 	log  log.Logger
 	hook hook.Service
@@ -63,20 +64,20 @@ func (t *h2cTransportWrapper) RoundTrip(req *http.Request) (*http.Response, erro
 
 func NewH2cRoundTripper(log log.Logger, hook hook.Service) http.RoundTripper {
 	return &h2cTransportWrapper{
-		httpTransport: &http.Transport{
+		httpTransport: otelhttp.NewTransport(&http.Transport{
 			DialContext: (&net.Dialer{
 				Timeout:   10 * time.Second,
 				KeepAlive: 1 * time.Minute,
 			}).DialContext,
 			DisableCompression: true,
-		},
-		grpcTransport: &http2.Transport{
+		}),
+		grpcTransport: otelhttp.NewTransport(&http2.Transport{
 			DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
 				return net.Dial(network, addr)
 			},
 			AllowHTTP:          true,
 			DisableCompression: true,
-		},
+		}),
 		log:  log,
 		hook: hook,
 	}

pkg/httputil/context.go

@@ -10,7 +10,6 @@ type (
 	contextRequestBodyKey struct{}
 	contextPathParamsKey  struct{}
 	routeKey              struct{}
-	varsKey               struct{}
 )
 
 func SetContextWithRequestBody(ctx context.Context, body []byte) context.Context {
@@ -31,10 +30,13 @@ func GetPathParamsFromContext(ctx context.Context) (map[string]string, bool) {
 	return params, ok
 }
 
-func SetContextWithMuxRouteAndVars(ctx context.Context, route *mux.Route, vars map[string]string) context.Context {
-	ctx = context.WithValue(ctx, routeKey{}, route)
-	if len(vars) > 0 {
-		ctx = context.WithValue(ctx, varsKey{}, vars)
+func SetContextWithMuxRoute(ctx context.Context, route *mux.Route) context.Context {
+	return context.WithValue(ctx, routeKey{}, route)
+}
+
+func GetMuxRoute(ctx context.Context) *mux.Route {
+	if rv := ctx.Value(routeKey{}); rv != nil {
+		return rv.(*mux.Route)
 	}
-	return ctx
+	return nil
 }

pkg/httputil/route.go

@@ -0,0 +1,22 @@
+package httputil
+
+import (
+	"net/http"
+
+	"github.com/gorilla/mux"
+)
+
+func MuxRouteName(r *http.Request) string {
+	route := mux.CurrentRoute(r)
+	if nil == route {
+		return "NotFoundHandler"
+	}
+	if n := route.GetName(); n != "" {
+		return n
+	}
+	if n, _ := route.GetPathTemplate(); n != "" {
+		return n
+	}
+	n, _ := route.GetHostTemplate()
+	return n
+}