forked from reddit/baseplate.go
/
thrift.go
155 lines (142 loc) · 4.26 KB
/
thrift.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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package tracing
import (
"context"
"fmt"
"strconv"
"github.com/apache/thrift/lib/go/thrift"
opentracing "github.com/opentracing/opentracing-go"
"github.com/fizx/baseplate.go/set"
"github.com/fizx/baseplate.go/thriftbp"
)
// StartSpanFromThriftContext creates a server span from thrift context object.
//
// This span would usually be used as the span of the whole thrift endpoint
// handler, and the parent of the child-spans.
//
// Caller should pass in the context object they got from thrift library,
// which would have all the required headers already injected.
//
// Please note that "Sampled" header is default to false according to baseplate
// spec, so if the context object doesn't have headers injected correctly,
// this span (and all its child-spans) will never be sampled,
// unless debug flag was set explicitly later.
//
// If any of the tracing related thrift header is present but malformed,
// it will be ignored.
// The error will also be logged if InitGlobalTracer was last called with a
// non-nil logger.
// Absent tracing related headers are always silently ignored.
func StartSpanFromThriftContext(ctx context.Context, name string) (context.Context, *Span) {
logger := globalTracer.getLogger()
span := newSpan(nil, name, SpanTypeServer)
defer func() {
onCreateServerSpan(span)
span.onStart()
}()
ctx = opentracing.ContextWithSpan(ctx, span)
if str, ok := thrift.GetHeader(ctx, thriftbp.HeaderTracingTrace); ok {
if id, err := strconv.ParseUint(str, 10, 64); err != nil {
if logger != nil {
logger(fmt.Sprintf(
"Malformed trace id in thrift ctx: %q, %v",
str,
err,
))
}
} else {
span.trace.traceID = id
}
}
if str, ok := thrift.GetHeader(ctx, thriftbp.HeaderTracingSpan); ok {
if id, err := strconv.ParseUint(str, 10, 64); err != nil {
if logger != nil {
logger(fmt.Sprintf(
"Malformed span id in thrift ctx: %q, %v",
str,
err,
))
}
} else {
span.trace.parentID = id
}
}
if str, ok := thrift.GetHeader(ctx, thriftbp.HeaderTracingFlags); ok {
if flags, err := strconv.ParseInt(str, 10, 64); err != nil {
if logger != nil {
logger(fmt.Sprintf(
"Malformed flags in thrift ctx: %q, %v",
str,
err,
))
}
} else {
span.trace.flags = flags
}
}
str, ok := thrift.GetHeader(ctx, thriftbp.HeaderTracingSampled)
sampled := ok && str == thriftbp.HeaderTracingSampledTrue
span.trace.sampled = sampled
return ctx, span
}
// CreateThriftContextFromSpan injects span info into a context object that can
// be used in thrift client code.
//
// Caller should first create a client child-span for the thrift call as usual,
// then use that span and the parent context object with this call,
// then use the returned context object in the thrift call.
// Something like:
//
// span, clientCtx := opentracing.StartSpanFromContext(
// ctx,
// "myCall",
// tracing.SpanTypeOption{Type: SpanTypeClient},
// )
// result, err := client.MyCall(clientCtx, arg1, arg2)
// // Or: span.Stop(ctx, err)
// span.FinishWithOptions(tracing.FinishOptions{
// Ctx: ctx,
// Err: err,
// }.Convert())
func CreateThriftContextFromSpan(ctx context.Context, span *Span) context.Context {
headers := set.StringSliceToSet(thrift.GetWriteHeaderList(ctx))
ctx = thrift.SetHeader(
ctx,
thriftbp.HeaderTracingTrace,
strconv.FormatUint(span.trace.traceID, 10),
)
headers.Add(thriftbp.HeaderTracingTrace)
ctx = thrift.SetHeader(
ctx,
thriftbp.HeaderTracingSpan,
strconv.FormatUint(span.trace.spanID, 10),
)
headers.Add(thriftbp.HeaderTracingSpan)
ctx = thrift.SetHeader(
ctx,
thriftbp.HeaderTracingFlags,
strconv.FormatInt(span.trace.flags, 10),
)
headers.Add(thriftbp.HeaderTracingFlags)
if span.trace.parentID != 0 {
ctx = thrift.SetHeader(
ctx,
thriftbp.HeaderTracingParent,
strconv.FormatUint(span.trace.parentID, 10),
)
headers.Add(thriftbp.HeaderTracingParent)
} else {
headers.Remove(thriftbp.HeaderTracingParent)
}
if span.trace.sampled {
ctx = thrift.SetHeader(
ctx,
thriftbp.HeaderTracingSampled,
thriftbp.HeaderTracingSampledTrue,
)
headers.Add(thriftbp.HeaderTracingSampled)
} else {
headers.Remove(thriftbp.HeaderTracingSampled)
}
ctx = thrift.SetWriteHeaderList(ctx, headers.ToSlice())
return ctx
}