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