-
Notifications
You must be signed in to change notification settings - Fork 1.1k
/
tracing.go
133 lines (111 loc) · 3.4 KB
/
tracing.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
package tracing
import (
"context"
"fmt"
"io"
"net/http"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/propagation"
"go.opentelemetry.io/otel/semconv"
"go.opentelemetry.io/otel/trace"
)
// StartSpan starts a new span as a child of the span in context.
// If there is no span in context then this is a no-op.
func StartSpan(ctx context.Context, operationName string, opts ...trace.SpanOption) (trace.Span, context.Context) {
parent := trace.SpanFromContext(ctx)
tracer := trace.NewNoopTracerProvider().Tracer("")
if parent != nil {
tracer = parent.Tracer()
}
ctx, span := tracer.Start(ctx, operationName, opts...)
return span, ctx
}
// FinishWithError finalizes the span and sets the error if one is passed
func FinishWithError(span trace.Span, err error) {
if err != nil {
span.RecordError(err)
if _, ok := err.(interface {
Cause() error
}); ok {
span.SetAttributes(attribute.String(string(semconv.ExceptionStacktraceKey), fmt.Sprintf("%+v", err)))
}
span.SetStatus(codes.Error, err.Error())
}
span.End()
}
// ContextWithSpanFromContext sets the tracing span of a context from other
// context if one is not already set. Alternative would be
// context.WithoutCancel() that would copy the context but reset ctx.Done
func ContextWithSpanFromContext(ctx, ctx2 context.Context) context.Context {
// if already is a span then noop
if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
return ctx
}
if span := trace.SpanFromContext(ctx2); span != nil {
return trace.ContextWithSpan(ctx, span)
}
return ctx
}
var DefaultTransport http.RoundTripper = &Transport{
RoundTripper: NewTransport(http.DefaultTransport),
}
var DefaultClient = &http.Client{
Transport: DefaultTransport,
}
var propagators = propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
type Transport struct {
http.RoundTripper
}
func NewTransport(rt http.RoundTripper) http.RoundTripper {
// TODO: switch to otelhttp. needs upstream updates to avoid transport-global tracer
return &Transport{
RoundTripper: rt,
}
}
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
span := trace.SpanFromContext(req.Context())
if span == nil { // no tracer connected with either request or transport
return t.RoundTripper.RoundTrip(req)
}
ctx, span := span.Tracer().Start(req.Context(), req.Method)
req = req.WithContext(ctx)
span.SetAttributes(semconv.HTTPClientAttributesFromHTTPRequest(req)...)
propagators.Inject(ctx, propagation.HeaderCarrier(req.Header))
resp, err := t.RoundTripper.RoundTrip(req)
if err != nil {
span.RecordError(err)
span.End()
return resp, err
}
span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(resp.StatusCode)...)
span.SetStatus(semconv.SpanStatusFromHTTPStatusCode(resp.StatusCode))
if req.Method == "HEAD" {
span.End()
} else {
resp.Body = &wrappedBody{ctx: ctx, span: span, body: resp.Body}
}
return resp, err
}
type wrappedBody struct {
ctx context.Context
span trace.Span
body io.ReadCloser
}
var _ io.ReadCloser = &wrappedBody{}
func (wb *wrappedBody) Read(b []byte) (int, error) {
n, err := wb.body.Read(b)
switch err {
case nil:
// nothing to do here but fall through to the return
case io.EOF:
wb.span.End()
default:
wb.span.RecordError(err)
}
return n, err
}
func (wb *wrappedBody) Close() error {
wb.span.End()
return wb.body.Close()
}