forked from DataDog/dd-trace-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
grpc.go
93 lines (86 loc) · 2.89 KB
/
grpc.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
//go:generate protoc -I . fixtures_test.proto --go_out=plugins=grpc:.
// Package grpc provides functions to trace the google.golang.org/grpc package v1.2.
package grpc // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/google.golang.org/grpc.v12"
import (
"net"
"gopkg.in/DataDog/dd-trace-go.v1/contrib/google.golang.org/internal/grpcutil"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
context "golang.org/x/net/context"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
)
// UnaryServerInterceptor will trace requests to the given grpc server.
func UnaryServerInterceptor(opts ...InterceptorOption) grpc.UnaryServerInterceptor {
cfg := new(interceptorConfig)
defaults(cfg)
for _, fn := range opts {
fn(cfg)
}
if cfg.serviceName == "" {
cfg.serviceName = "grpc.server"
}
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
span, ctx := startSpanFromContext(ctx, info.FullMethod, cfg.serviceName)
resp, err := handler(ctx, req)
span.Finish(tracer.WithError(err))
return resp, err
}
}
func startSpanFromContext(ctx context.Context, method, service string) (ddtrace.Span, context.Context) {
opts := []ddtrace.StartSpanOption{
tracer.ServiceName(service),
tracer.ResourceName(method),
tracer.Tag(tagMethod, method),
tracer.SpanType(ext.AppTypeRPC),
}
md, _ := metadata.FromContext(ctx) // nil is ok
if sctx, err := tracer.Extract(grpcutil.MDCarrier(md)); err == nil {
opts = append(opts, tracer.ChildOf(sctx))
}
return tracer.StartSpanFromContext(ctx, "grpc.server", opts...)
}
// UnaryClientInterceptor will add tracing to a gprc client.
func UnaryClientInterceptor(opts ...InterceptorOption) grpc.UnaryClientInterceptor {
cfg := new(interceptorConfig)
defaults(cfg)
for _, fn := range opts {
fn(cfg)
}
if cfg.serviceName == "" {
cfg.serviceName = "grpc.client"
}
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
var (
span ddtrace.Span
p peer.Peer
)
span, ctx = tracer.StartSpanFromContext(ctx, "grpc.client",
tracer.Tag(tagMethod, method),
tracer.SpanType(ext.AppTypeRPC),
)
md, ok := metadata.FromContext(ctx)
if !ok {
md = metadata.MD{}
}
_ = tracer.Inject(span.Context(), grpcutil.MDCarrier(md))
ctx = metadata.NewContext(ctx, md)
opts = append(opts, grpc.Peer(&p))
err := invoker(ctx, method, req, reply, cc, opts...)
if p.Addr != nil {
addr := p.Addr.String()
host, port, err := net.SplitHostPort(addr)
if err == nil {
if host != "" {
span.SetTag(ext.TargetHost, host)
}
span.SetTag(ext.TargetPort, port)
}
}
span.SetTag(tagCode, grpc.Code(err).String())
span.Finish(tracer.WithError(err))
return err
}
}