1
// Copyright (c) 2023 Uber Technologies, Inc.
2
//
3
// Permission is hereby granted, free of charge, to any person obtaining a copy
4
// of this software and associated documentation files (the "Software"), to deal
5
// in the Software without restriction, including without limitation the rights
6
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7
// copies of the Software, and to permit persons to whom the Software is
8
// furnished to do so, subject to the following conditions:
9
//
10
// The above copyright notice and this permission notice shall be included in
11
// all copies or substantial portions of the Software.
12
//
13
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
19
// THE SOFTWARE.
20

21
// Package stacktrace provides support for gathering stack traces
22
// efficiently.
23
package stacktrace
24

25
import (
26
	"runtime"
27

28
	"go.uber.org/zap/buffer"
29
	"go.uber.org/zap/internal/bufferpool"
30
	"go.uber.org/zap/internal/pool"
31
)
32

33
var _stackPool = pool.New(func() *Stack {
34
	return &Stack{
35
		storage: make([]uintptr, 64),
36
	}
37
})
38

39
// Stack is a captured stack trace.
40
type Stack struct {
41
	pcs    []uintptr // program counters; always a subslice of storage
42
	frames *runtime.Frames
43

44
	// The size of pcs varies depending on requirements:
45
	// it will be one if the only the first frame was requested,
46
	// and otherwise it will reflect the depth of the call stack.
47
	//
48
	// storage decouples the slice we need (pcs) from the slice we pool.
49
	// We will always allocate a reasonably large storage, but we'll use
50
	// only as much of it as we need.
51
	storage []uintptr
52
}
53

54
// Depth specifies how deep of a stack trace should be captured.
55
type Depth int
56

57
const (
58
	// First captures only the first frame.
59
	First Depth = iota
60

61
	// Full captures the entire call stack, allocating more
62
	// storage for it if needed.
63
	Full
64
)
65

66
// Capture captures a stack trace of the specified depth, skipping
67
// the provided number of frames. skip=0 identifies the caller of
68
// Capture.
69
//
70
// The caller must call Free on the returned stacktrace after using it.
71
func Capture(skip int, depth Depth) *Stack {
72
	stack := _stackPool.Get()
73

74
	switch depth {
75
	case First:
76
		stack.pcs = stack.storage[:1]
77
	case Full:
78
		stack.pcs = stack.storage
79
	}
80

81
	// Unlike other "skip"-based APIs, skip=0 identifies runtime.Callers
82
	// itself. +2 to skip captureStacktrace and runtime.Callers.
83
	numFrames := runtime.Callers(
84
		skip+2,
85
		stack.pcs,
86
	)
87

88
	// runtime.Callers truncates the recorded stacktrace if there is no
89
	// room in the provided slice. For the full stack trace, keep expanding
90
	// storage until there are fewer frames than there is room.
91
	if depth == Full {
92
		pcs := stack.pcs
93
		for numFrames == len(pcs) {
94
			pcs = make([]uintptr, len(pcs)*2)
95
			numFrames = runtime.Callers(skip+2, pcs)
96
		}
97

98
		// Discard old storage instead of returning it to the pool.
99
		// This will adjust the pool size over time if stack traces are
100
		// consistently very deep.
101
		stack.storage = pcs
102
		stack.pcs = pcs[:numFrames]
103
	} else {
104
		stack.pcs = stack.pcs[:numFrames]
105
	}
106

107
	stack.frames = runtime.CallersFrames(stack.pcs)
108
	return stack
109
}
110

111
// Free releases resources associated with this stacktrace
112
// and returns it back to the pool.
113
func (st *Stack) Free() {
114
	st.frames = nil
115
	st.pcs = nil
116
	_stackPool.Put(st)
117
}
118

119
// Count reports the total number of frames in this stacktrace.
120
// Count DOES NOT change as Next is called.
121
func (st *Stack) Count() int {
122
	return len(st.pcs)
123
}
124

125
// Next returns the next frame in the stack trace,
126
// and a boolean indicating whether there are more after it.
127
func (st *Stack) Next() (_ runtime.Frame, more bool) {
128
	return st.frames.Next()
129
}
130

131
// Take returns a string representation of the current stacktrace.
132
//
133
// skip is the number of frames to skip before recording the stack trace.
134
// skip=0 identifies the caller of Take.
135
func Take(skip int) string {
136
	stack := Capture(skip+1, Full)
137
	defer stack.Free()
138

139
	buffer := bufferpool.Get()
140
	defer buffer.Free()
141

142
	stackfmt := NewFormatter(buffer)
143
	stackfmt.FormatStack(stack)
144
	return buffer.String()
145
}
146

147
// Formatter formats a stack trace into a readable string representation.
148
type Formatter struct {
149
	b        *buffer.Buffer
150
	nonEmpty bool // whehther we've written at least one frame already
151
}
152

153
// NewFormatter builds a new Formatter.
154
func NewFormatter(b *buffer.Buffer) Formatter {
155
	return Formatter{b: b}
156
}
157

158
// FormatStack formats all remaining frames in the provided stacktrace -- minus
159
// the final runtime.main/runtime.goexit frame.
160
func (sf *Formatter) FormatStack(stack *Stack) {
161
	// Note: On the last iteration, frames.Next() returns false, with a valid
162
	// frame, but we ignore this frame. The last frame is a runtime frame which
163
	// adds noise, since it's only either runtime.main or runtime.goexit.
164
	for frame, more := stack.Next(); more; frame, more = stack.Next() {
165
		sf.FormatFrame(frame)
166
	}
167
}
168

169
// FormatFrame formats the given frame.
170
func (sf *Formatter) FormatFrame(frame runtime.Frame) {
171
	if sf.nonEmpty {
172
		sf.b.AppendByte('\n')
173
	}
174
	sf.nonEmpty = true
175
	sf.b.AppendString(frame.Function)
176
	sf.b.AppendByte('\n')
177
	sf.b.AppendByte('\t')
178
	sf.b.AppendString(frame.File)
179
	sf.b.AppendByte(':')
180
	sf.b.AppendInt(int64(frame.Line))
181
}
182

183