1// Copyright 2019 The etcd Authors
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
15// Package traceutil implements tracing utilities using "context".
16package traceutil
17
18import (
19	"bytes"
20	"context"
21	"fmt"
22	"math/rand"
23	"time"
24
25	"go.uber.org/zap"
26)
27
28const (
29	TraceKey     = "trace"
30	StartTimeKey = "startTime"
31)
32
33// Field is a kv pair to record additional details of the trace.
34type Field struct {
35	Key   string
36	Value interface{}
37}
38
39func (f *Field) format() string {
40	return fmt.Sprintf("%s:%v; ", f.Key, f.Value)
41}
42
43func writeFields(fields []Field) string {
44	if len(fields) == 0 {
45		return ""
46	}
47	var buf bytes.Buffer
48	buf.WriteString("{")
49	for _, f := range fields {
50		buf.WriteString(f.format())
51	}
52	buf.WriteString("}")
53	return buf.String()
54}
55
56type Trace struct {
57	operation    string
58	lg           *zap.Logger
59	fields       []Field
60	startTime    time.Time
61	steps        []step
62	stepDisabled bool
63	isEmpty      bool
64}
65
66type step struct {
67	time            time.Time
68	msg             string
69	fields          []Field
70	isSubTraceStart bool
71	isSubTraceEnd   bool
72}
73
74func New(op string, lg *zap.Logger, fields ...Field) *Trace {
75	return &Trace{operation: op, lg: lg, startTime: time.Now(), fields: fields}
76}
77
78// TODO returns a non-nil, empty Trace
79func TODO() *Trace {
80	return &Trace{isEmpty: true}
81}
82
83func Get(ctx context.Context) *Trace {
84	if trace, ok := ctx.Value(TraceKey).(*Trace); ok && trace != nil {
85		return trace
86	}
87	return TODO()
88}
89
90func (t *Trace) GetStartTime() time.Time {
91	return t.startTime
92}
93
94func (t *Trace) SetStartTime(time time.Time) {
95	t.startTime = time
96}
97
98func (t *Trace) InsertStep(at int, time time.Time, msg string, fields ...Field) {
99	newStep := step{time: time, msg: msg, fields: fields}
100	if at < len(t.steps) {
101		t.steps = append(t.steps[:at+1], t.steps[at:]...)
102		t.steps[at] = newStep
103	} else {
104		t.steps = append(t.steps, newStep)
105	}
106}
107
108// StartSubTrace adds step to trace as a start sign of sublevel trace
109// All steps in the subtrace will log out the input fields of this function
110func (t *Trace) StartSubTrace(fields ...Field) {
111	t.steps = append(t.steps, step{fields: fields, isSubTraceStart: true})
112}
113
114// StopSubTrace adds step to trace as a end sign of sublevel trace
115// All steps in the subtrace will log out the input fields of this function
116func (t *Trace) StopSubTrace(fields ...Field) {
117	t.steps = append(t.steps, step{fields: fields, isSubTraceEnd: true})
118}
119
120// Step adds step to trace
121func (t *Trace) Step(msg string, fields ...Field) {
122	if !t.stepDisabled {
123		t.steps = append(t.steps, step{time: time.Now(), msg: msg, fields: fields})
124	}
125}
126
127// StepWithFunction will measure the input function as a single step
128func (t *Trace) StepWithFunction(f func(), msg string, fields ...Field) {
129	t.disableStep()
130	f()
131	t.enableStep()
132	t.Step(msg, fields...)
133}
134
135func (t *Trace) AddField(fields ...Field) {
136	for _, f := range fields {
137		if !t.updateFieldIfExist(f) {
138			t.fields = append(t.fields, f)
139		}
140	}
141}
142
143func (t *Trace) IsEmpty() bool {
144	return t.isEmpty
145}
146
147// Log dumps all steps in the Trace
148func (t *Trace) Log() {
149	t.LogWithStepThreshold(0)
150}
151
152// LogIfLong dumps logs if the duration is longer than threshold
153func (t *Trace) LogIfLong(threshold time.Duration) {
154	if time.Since(t.startTime) > threshold {
155		stepThreshold := threshold / time.Duration(len(t.steps)+1)
156		t.LogWithStepThreshold(stepThreshold)
157	}
158}
159
160// LogAllStepsIfLong dumps all logs if the duration is longer than threshold
161func (t *Trace) LogAllStepsIfLong(threshold time.Duration) {
162	if time.Since(t.startTime) > threshold {
163		t.LogWithStepThreshold(0)
164	}
165}
166
167// LogWithStepThreshold only dumps step whose duration is longer than step threshold
168func (t *Trace) LogWithStepThreshold(threshold time.Duration) {
169	msg, fs := t.logInfo(threshold)
170	if t.lg != nil {
171		t.lg.Info(msg, fs...)
172	}
173}
174
175func (t *Trace) logInfo(threshold time.Duration) (string, []zap.Field) {
176	endTime := time.Now()
177	totalDuration := endTime.Sub(t.startTime)
178	traceNum := rand.Int31()
179	msg := fmt.Sprintf("trace[%d] %s", traceNum, t.operation)
180
181	var steps []string
182	lastStepTime := t.startTime
183	for i := 0; i < len(t.steps); i++ {
184		step := t.steps[i]
185		// add subtrace common fields which defined at the beginning to each sub-steps
186		if step.isSubTraceStart {
187			for j := i + 1; j < len(t.steps) && !t.steps[j].isSubTraceEnd; j++ {
188				t.steps[j].fields = append(step.fields, t.steps[j].fields...)
189			}
190			continue
191		}
192		// add subtrace common fields which defined at the end to each sub-steps
193		if step.isSubTraceEnd {
194			for j := i - 1; j >= 0 && !t.steps[j].isSubTraceStart; j-- {
195				t.steps[j].fields = append(step.fields, t.steps[j].fields...)
196			}
197			continue
198		}
199	}
200	for i := 0; i < len(t.steps); i++ {
201		step := t.steps[i]
202		if step.isSubTraceStart || step.isSubTraceEnd {
203			continue
204		}
205		stepDuration := step.time.Sub(lastStepTime)
206		if stepDuration > threshold {
207			steps = append(steps, fmt.Sprintf("trace[%d] '%v' %s (duration: %v)",
208				traceNum, step.msg, writeFields(step.fields), stepDuration))
209		}
210		lastStepTime = step.time
211	}
212
213	fs := []zap.Field{zap.String("detail", writeFields(t.fields)),
214		zap.Duration("duration", totalDuration),
215		zap.Time("start", t.startTime),
216		zap.Time("end", endTime),
217		zap.Strings("steps", steps),
218		zap.Int("step_count", len(steps))}
219	return msg, fs
220}
221
222func (t *Trace) updateFieldIfExist(f Field) bool {
223	for i, v := range t.fields {
224		if v.Key == f.Key {
225			t.fields[i].Value = f.Value
226			return true
227		}
228	}
229	return false
230}
231
232// disableStep sets the flag to prevent the trace from adding steps
233func (t *Trace) disableStep() {
234	t.stepDisabled = true
235}
236
237// enableStep re-enable the trace to add steps
238func (t *Trace) enableStep() {
239	t.stepDisabled = false
240}
241