/
reporter.go
64 lines (54 loc) · 1.54 KB
/
reporter.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
// Copyright (c) The go-grpc-middleware Authors.
// Licensed under the Apache License 2.0.
package tracing
import (
"context"
"io"
"time"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
)
type reporter struct {
ctx context.Context
span Span
receivedMessageID int
sentMessageID int
}
func (o *reporter) PostCall(err error, _ time.Duration) {
// Finish span.
if err != nil && err != io.EOF {
s, _ := status.FromError(err)
o.span.SetStatus(s.Code(), s.Message())
o.span.SetAttributes(statusCodeAttr(s.Code()))
} else {
o.span.SetAttributes(statusCodeAttr(codes.OK))
}
o.span.End()
}
func (o *reporter) PostMsgSend(payload interface{}, err error, d time.Duration) {
o.sentMessageID++
addEvent(o.span, RPCMessageTypeSent, o.sentMessageID, payload)
}
func (o *reporter) PostMsgReceive(payload interface{}, err error, d time.Duration) {
o.receivedMessageID++
addEvent(o.span, RPCMessageTypeReceived, o.receivedMessageID, payload)
}
func addEvent(span Span, messageType string, messageID int, payload interface{}) {
if p, ok := payload.(proto.Message); ok {
span.AddEvent("message",
rpcMessageTypeKey, messageType,
rpcMessageIDKey, messageID,
rpcMessageUncompressedSizeKey, proto.Size(p),
)
return
}
span.AddEvent("message",
rpcMessageTypeKey, messageType,
rpcMessageIDKey, messageID,
)
}
// statusCodeAttr returns status code attribute based on given gRPC code
func statusCodeAttr(c codes.Code) []interface{} {
return []interface{}{grpcStatusCodeKey, int64(c)}
}