/
unary_client.go
54 lines (45 loc) · 1.47 KB
/
unary_client.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
package tracing
import (
"context"
"fmt"
"github.com/opentracing/opentracing-go"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
// WithUnaryClient 客户端普通方法调用追踪
func WithUnaryClient(opts ...Option) grpc.DialOption {
var defaultOptions = &options{
tracer: opentracing.GlobalTracer(),
payloadMarshal: defaultPayloadMarshal,
opName: defaultOperationName,
}
defaultOptions = mergeOptions(defaultOptions, opts)
return grpc.WithChainUnaryInterceptor(unaryClientTracing(defaultOptions))
}
func unaryClientTracing(defaultOptions *options) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
var grpcOpts, nOpts = filterOptions(opts)
var opt = mergeOptions(defaultOptions, nOpts)
if opt.disable {
return invoker(ctx, method, req, reply, cc, grpcOpts...)
}
var opName = opt.opName(ctx, method)
var nCtx, nSpan, err = clientSpanFromContext(ctx, opt.tracer, fmt.Sprintf("[GRPC Client] %s", opName))
if err != nil {
return err
}
if opt.payload {
var md, _ = metadata.FromOutgoingContext(ctx)
traceHeader(nSpan, md)
nSpan.LogKV("Req", opt.payloadMarshal(req))
}
err = invoker(nCtx, method, req, reply, cc, grpcOpts...)
if opt.payload {
if err == nil && reply != nil {
nSpan.LogKV("Recv", opt.payloadMarshal(reply))
}
}
finish(nSpan, err)
return err
}
}