1package opencensus
2
3// Copyright 2018 Microsoft Corporation
4//
5//  Licensed under the Apache License, Version 2.0 (the "License");
6//  you may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at
8//
9//      http://www.apache.org/licenses/LICENSE-2.0
10//
11//  Unless required by applicable law or agreed to in writing, software
12//  distributed under the License is distributed on an "AS IS" BASIS,
13//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14//  See the License for the specific language governing permissions and
15//  limitations under the License.
16
17import (
18	"context"
19	"fmt"
20	"net/http"
21	"os"
22
23	"contrib.go.opencensus.io/exporter/ocagent"
24	"github.com/Azure/go-autorest/tracing"
25	"go.opencensus.io/plugin/ochttp"
26	"go.opencensus.io/plugin/ochttp/propagation/tracecontext"
27	"go.opencensus.io/stats/view"
28	"go.opencensus.io/trace"
29)
30
31func init() {
32	enableFromEnv()
33}
34
35// split out for testing purposes
36func enableFromEnv() {
37	if _, ok := os.LookupEnv("AZURE_SDK_TRACING_ENABLED"); ok {
38		agentEndpoint, ok := os.LookupEnv("OCAGENT_TRACE_EXPORTER_ENDPOINT")
39		if ok {
40			EnableWithAIForwarding(agentEndpoint)
41		} else {
42			Enable()
43		}
44	}
45}
46
47var defaultTracer = newocTracer()
48
49type ocTracer struct {
50	// Sampler is the tracing sampler. If tracing is disabled it will never sample. Otherwise
51	// it will be using the parent sampler or the default.
52	sampler trace.Sampler
53
54	// Views for metric instrumentation.
55	views map[string]*view.View
56
57	// the trace exporter
58	traceExporter trace.Exporter
59}
60
61func newocTracer() *ocTracer {
62	return &ocTracer{
63		sampler: trace.NeverSample(),
64		views:   map[string]*view.View{},
65	}
66}
67
68// NewTransport returns a new instance of a tracing-aware RoundTripper.
69func (oct ocTracer) NewTransport(base *http.Transport) http.RoundTripper {
70	return &ochttp.Transport{
71		Base:        base,
72		Propagation: &tracecontext.HTTPFormat{},
73		GetStartOptions: func(*http.Request) trace.StartOptions {
74			return trace.StartOptions{
75				Sampler: oct.sampler,
76			}
77		},
78	}
79}
80
81// StartSpan starts a trace span
82func (oct ocTracer) StartSpan(ctx context.Context, name string) context.Context {
83	ctx, _ = trace.StartSpan(ctx, name, trace.WithSampler(oct.sampler))
84	return ctx
85}
86
87// EndSpan ends a previously started span stored in the context
88func (oct ocTracer) EndSpan(ctx context.Context, httpStatusCode int, err error) {
89	span := trace.FromContext(ctx)
90
91	if span == nil {
92		return
93	}
94
95	if err != nil {
96		span.SetStatus(trace.Status{Message: err.Error(), Code: toTraceStatusCode(httpStatusCode)})
97	}
98	span.End()
99}
100
101// Enable will start instrumentation for metrics and traces.
102func Enable() error {
103	defaultTracer.sampler = nil
104
105	// register the views for HTTP metrics
106	clientViews := []*view.View{
107		ochttp.ClientCompletedCount,
108		ochttp.ClientRoundtripLatencyDistribution,
109		ochttp.ClientReceivedBytesDistribution,
110		ochttp.ClientSentBytesDistribution,
111	}
112	for _, cv := range clientViews {
113		vn := fmt.Sprintf("Azure/go-autorest/tracing/opencensus-%s", cv.Name)
114		defaultTracer.views[vn] = cv.WithName(vn)
115		err := view.Register(defaultTracer.views[vn])
116		if err != nil {
117			return err
118		}
119	}
120	tracing.Register(defaultTracer)
121	return nil
122}
123
124// Disable will disable instrumentation for metrics and traces.
125func Disable() {
126	// unregister any previously registered metrics
127	for _, v := range defaultTracer.views {
128		view.Unregister(v)
129	}
130	defaultTracer.sampler = trace.NeverSample()
131	if defaultTracer.traceExporter != nil {
132		trace.UnregisterExporter(defaultTracer.traceExporter)
133	}
134	tracing.Register(nil)
135}
136
137// EnableWithAIForwarding will start instrumentation and will connect to app insights forwarder
138// exporter making the metrics and traces available in app insights.
139func EnableWithAIForwarding(agentEndpoint string) error {
140	err := Enable()
141	if err != nil {
142		return err
143	}
144
145	defaultTracer.traceExporter, err = ocagent.NewExporter(ocagent.WithInsecure(), ocagent.WithAddress(agentEndpoint))
146	if err != nil {
147		return err
148	}
149	trace.RegisterExporter(defaultTracer.traceExporter)
150	return nil
151}
152
153// toTraceStatusCode converts HTTP Codes to OpenCensus codes as defined
154// at https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/HTTP.md#status
155func toTraceStatusCode(httpStatusCode int) int32 {
156	switch {
157	case http.StatusOK <= httpStatusCode && httpStatusCode < http.StatusBadRequest:
158		return trace.StatusCodeOK
159	case httpStatusCode == http.StatusBadRequest:
160		return trace.StatusCodeInvalidArgument
161	case httpStatusCode == http.StatusUnauthorized: // 401 is actually unauthenticated.
162		return trace.StatusCodeUnauthenticated
163	case httpStatusCode == http.StatusForbidden:
164		return trace.StatusCodePermissionDenied
165	case httpStatusCode == http.StatusNotFound:
166		return trace.StatusCodeNotFound
167	case httpStatusCode == http.StatusTooManyRequests:
168		return trace.StatusCodeResourceExhausted
169	case httpStatusCode == 499:
170		return trace.StatusCodeCancelled
171	case httpStatusCode == http.StatusNotImplemented:
172		return trace.StatusCodeUnimplemented
173	case httpStatusCode == http.StatusServiceUnavailable:
174		return trace.StatusCodeUnavailable
175	case httpStatusCode == http.StatusGatewayTimeout:
176		return trace.StatusCodeDeadlineExceeded
177	default:
178		return trace.StatusCodeUnknown
179	}
180}
181