Skip to content

Commit

Permalink
feat: enrich trace span information
Browse files Browse the repository at this point in the history
  • Loading branch information
mabdh committed Jun 12, 2024
1 parent d00bcf0 commit ef75c70
Show file tree
Hide file tree
Showing 10 changed files with 142 additions and 37 deletions.
11 changes: 9 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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>

Expand Down Expand Up @@ -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
```
Expand Down
19 changes: 3 additions & 16 deletions cmd/serve_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
}
4 changes: 2 additions & 2 deletions internal/proxy/middleware/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
24 changes: 19 additions & 5 deletions internal/proxy/middleware/observability/observability.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
}), ""),
}
}

Expand All @@ -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 {
Expand Down
42 changes: 42 additions & 0 deletions internal/proxy/middleware/otelpostprocessor/otelpostprocessor.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
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(routeName)

ctx := trace.ContextWithSpan(r.Context(), span)
r = r.WithContext(ctx)

m.next.ServeHTTP(rw, r)
}
Original file line number Diff line number Diff line change
@@ -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)
}
2 changes: 1 addition & 1 deletion internal/proxy/middleware/rulematch/route_matcher.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
13 changes: 7 additions & 6 deletions internal/proxy/roundtripper.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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,
}
Expand Down
13 changes: 8 additions & 5 deletions pkg/httputil/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,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
}
20 changes: 20 additions & 0 deletions pkg/httputil/route.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
package httputil

import (
"net/http"
)

func MuxRouteName(r *http.Request) string {
route := GetMuxRoute(r.Context())
if nil == route {
return "NotFoundHandler"
}
if n := route.GetName(); n != "" {
return n
}
if n, _ := route.GetPathTemplate(); n != "" {
return r.Method + " " + n
}
n, _ := route.GetHostTemplate()
return r.Method + " " + n
}

0 comments on commit ef75c70

Please sign in to comment.