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	"golang.org/x/xerrors"
23	"google.golang.org/api/googleapi"
24	"google.golang.org/genproto/googleapis/rpc/code"
25	"google.golang.org/grpc/status"
26)
27
28// StartSpan adds a span to the trace with the given name.
29func StartSpan(ctx context.Context, name string) context.Context {
30	ctx, _ = trace.StartSpan(ctx, name)
31	return ctx
32}
33
34// EndSpan ends a span with the given error.
35func EndSpan(ctx context.Context, err error) {
36	span := trace.FromContext(ctx)
37	if err != nil {
38		span.SetStatus(toStatus(err))
39	}
40	span.End()
41}
42
43// toStatus interrogates an error and converts it to an appropriate
44// OpenCensus status.
45func toStatus(err error) trace.Status {
46	var err2 *googleapi.Error
47	if ok := xerrors.As(err, &err2); ok {
48		return trace.Status{Code: httpStatusCodeToOCCode(err2.Code), Message: err2.Message}
49	} else if s, ok := status.FromError(err); ok {
50		return trace.Status{Code: int32(s.Code()), Message: s.Message()}
51	} else {
52		return trace.Status{Code: int32(code.Code_UNKNOWN), Message: err.Error()}
53	}
54}
55
56// TODO(deklerk): switch to using OpenCensus function when it becomes available.
57// Reference: https://github.com/googleapis/googleapis/blob/26b634d2724ac5dd30ae0b0cbfb01f07f2e4050e/google/rpc/code.proto
58func httpStatusCodeToOCCode(httpStatusCode int) int32 {
59	switch httpStatusCode {
60	case 200:
61		return int32(code.Code_OK)
62	case 499:
63		return int32(code.Code_CANCELLED)
64	case 500:
65		return int32(code.Code_UNKNOWN) // Could also be Code_INTERNAL, Code_DATA_LOSS
66	case 400:
67		return int32(code.Code_INVALID_ARGUMENT) // Could also be Code_OUT_OF_RANGE
68	case 504:
69		return int32(code.Code_DEADLINE_EXCEEDED)
70	case 404:
71		return int32(code.Code_NOT_FOUND)
72	case 409:
73		return int32(code.Code_ALREADY_EXISTS) // Could also be Code_ABORTED
74	case 403:
75		return int32(code.Code_PERMISSION_DENIED)
76	case 401:
77		return int32(code.Code_UNAUTHENTICATED)
78	case 429:
79		return int32(code.Code_RESOURCE_EXHAUSTED)
80	case 501:
81		return int32(code.Code_UNIMPLEMENTED)
82	case 503:
83		return int32(code.Code_UNAVAILABLE)
84	default:
85		return int32(code.Code_UNKNOWN)
86	}
87}
88
89// TODO: (odeke-em): perhaps just pass around spans due to the cost
90// incurred from using trace.FromContext(ctx) yet we could avoid
91// throwing away the work done by ctx, span := trace.StartSpan.
92func TracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
93	var attrs []trace.Attribute
94	for k, v := range attrMap {
95		var a trace.Attribute
96		switch v := v.(type) {
97		case string:
98			a = trace.StringAttribute(k, v)
99		case bool:
100			a = trace.BoolAttribute(k, v)
101		case int:
102			a = trace.Int64Attribute(k, int64(v))
103		case int64:
104			a = trace.Int64Attribute(k, v)
105		default:
106			a = trace.StringAttribute(k, fmt.Sprintf("%#v", v))
107		}
108		attrs = append(attrs, a)
109	}
110	trace.FromContext(ctx).Annotatef(attrs, format, args...)
111}
112