/
opentracing.go
92 lines (79 loc) · 2.41 KB
/
opentracing.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
package middleware
import (
"context"
"strings"
logg "github.com/oarkflow/log"
"github.com/opentracing/opentracing-go"
"github.com/opentracing/opentracing-go/ext"
"github.com/opentracing/opentracing-go/log"
"github.com/pkg/errors"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
type MDReaderWriter struct {
metadata.MD
}
func (c MDReaderWriter) ForeachKey(handler func(key, val string) error) error {
for k, vs := range c.MD {
for _, v := range vs {
if err := handler(k, v); err != nil {
return err
}
}
}
return nil
}
func (c MDReaderWriter) Set(key, val string) {
key = strings.ToLower(key)
c.MD[key] = append(c.MD[key], val)
}
func OpentracingClient(tracer opentracing.Tracer, parentCtx context.Context) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
span, _ := opentracing.StartSpanFromContext(parentCtx,
"call gRPC",
opentracing.Tag{Key: string(ext.Component), Value: "gRPC"},
ext.SpanKindRPCClient)
defer span.Finish()
md, ok := metadata.FromOutgoingContext(ctx)
if !ok {
md = metadata.New(nil)
} else {
md = md.Copy()
}
err := tracer.Inject(span.Context(), opentracing.TextMap, MDReaderWriter{md})
if err != nil {
span.LogFields(log.Error(errors.WithMessage(err, "inject-error")))
}
newCtx := metadata.NewOutgoingContext(ctx, md)
err = invoker(newCtx, method, req, reply, cc, opts...)
if err != nil {
span.LogFields(log.Error(errors.WithMessage(err, "call-error")))
}
return err
}
}
func OpentracingServer(tracer opentracing.Tracer) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (
resp any, err error) {
var parentCtx context.Context
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
md = metadata.New(nil)
}
spanContext, err := tracer.Extract(opentracing.TextMap, MDReaderWriter{md})
if err != nil && err != opentracing.ErrSpanContextNotFound {
logg.Error().Err(err).Msg("extract from metadata error")
} else {
span := tracer.StartSpan(
info.FullMethod,
ext.RPCServerOption(spanContext),
opentracing.Tag{Key: string(ext.Component), Value: "gRPC"},
ext.SpanKindRPCServer,
)
defer span.Finish()
parentCtx = opentracing.ContextWithSpan(ctx, span)
}
return handler(parentCtx, req)
}
}