-
Notifications
You must be signed in to change notification settings - Fork 3
/
intercept.go
55 lines (42 loc) · 1.31 KB
/
intercept.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
package etcdv3
import (
"context"
"github.com/abulo/ratel/v3/core/trace"
"google.golang.org/grpc"
)
func traceUnaryClientInterceptor() grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
span, ctx := trace.StartSpanFromContext(
ctx,
method,
trace.FromIncomingContext(ctx),
trace.TagComponent("gRPC"),
trace.TagSpanKind("server.unary"),
)
defer span.Finish()
err = invoker(ctx, method, req, reply, cc, opts...)
// span.SetTag("code", ecode.OK.Code)
// if err != nil {
// span.SetTag("code", ecode.ExtractCodes())
// }
return err
}
}
func traceStreamClientInterceptor() grpc.StreamClientInterceptor {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
span, ctx := trace.StartSpanFromContext(
ctx,
method,
trace.FromIncomingContext(ctx),
trace.TagComponent("gRPC"),
trace.TagSpanKind("server.unary"),
)
defer span.Finish()
clientStream, err := streamer(ctx, desc, cc, method, opts...)
// span.SetTag("code", codes.Ok)
// if err != nil {
// span.SetTag("code", codes.Error)
// }
return clientStream, err
}
}