1// Copyright 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package trace
16
17import (
18	"context"
19	"fmt"
20
21	"go.opencensus.io/trace"
22	"google.golang.org/api/googleapi"
23	"google.golang.org/genproto/googleapis/rpc/code"
24	"google.golang.org/grpc/status"
25)
26
27// StartSpan adds a span to the trace with the given name.
28func StartSpan(ctx context.Context, name string) context.Context {
29	ctx, _ = trace.StartSpan(ctx, name)
30	return ctx
31}
32
33// EndSpan ends a span with the given error.
34func EndSpan(ctx context.Context, err error) {
35	span := trace.FromContext(ctx)
36	if err != nil {
37		span.SetStatus(toStatus(err))
38	}
39	span.End()
40}
41
42// toStatus interrogates an error and converts it to an appropriate
43// OpenCensus status.
44func toStatus(err error) trace.Status {
45	if err2, ok := err.(*googleapi.Error); ok {
46		return trace.Status{Code: httpStatusCodeToOCCode(err2.Code), Message: err2.Message}
47	} else if s, ok := status.FromError(err); ok {
48		return trace.Status{Code: int32(s.Code()), Message: s.Message()}
49	} else {
50		return trace.Status{Code: int32(code.Code_UNKNOWN), Message: err.Error()}
51	}
52}
53
54// TODO(deklerk): switch to using OpenCensus function when it becomes available.
55// Reference: https://github.com/googleapis/googleapis/blob/26b634d2724ac5dd30ae0b0cbfb01f07f2e4050e/google/rpc/code.proto
56func httpStatusCodeToOCCode(httpStatusCode int) int32 {
57	switch httpStatusCode {
58	case 200:
59		return int32(code.Code_OK)
60	case 499:
61		return int32(code.Code_CANCELLED)
62	case 500:
63		return int32(code.Code_UNKNOWN) // Could also be Code_INTERNAL, Code_DATA_LOSS
64	case 400:
65		return int32(code.Code_INVALID_ARGUMENT) // Could also be Code_OUT_OF_RANGE
66	case 504:
67		return int32(code.Code_DEADLINE_EXCEEDED)
68	case 404:
69		return int32(code.Code_NOT_FOUND)
70	case 409:
71		return int32(code.Code_ALREADY_EXISTS) // Could also be Code_ABORTED
72	case 403:
73		return int32(code.Code_PERMISSION_DENIED)
74	case 401:
75		return int32(code.Code_UNAUTHENTICATED)
76	case 429:
77		return int32(code.Code_RESOURCE_EXHAUSTED)
78	case 501:
79		return int32(code.Code_UNIMPLEMENTED)
80	case 503:
81		return int32(code.Code_UNAVAILABLE)
82	default:
83		return int32(code.Code_UNKNOWN)
84	}
85}
86
87// TODO: (odeke-em): perhaps just pass around spans due to the cost
88// incurred from using trace.FromContext(ctx) yet we could avoid
89// throwing away the work done by ctx, span := trace.StartSpan.
90func TracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
91	var attrs []trace.Attribute
92	for k, v := range attrMap {
93		var a trace.Attribute
94		switch v := v.(type) {
95		case string:
96			a = trace.StringAttribute(k, v)
97		case bool:
98			a = trace.BoolAttribute(k, v)
99		case int:
100			a = trace.Int64Attribute(k, int64(v))
101		case int64:
102			a = trace.Int64Attribute(k, v)
103		default:
104			a = trace.StringAttribute(k, fmt.Sprintf("%#v", v))
105		}
106		attrs = append(attrs, a)
107	}
108	trace.FromContext(ctx).Annotatef(attrs, format, args...)
109}
110