95 lines
2.5 KiB
Go
95 lines
2.5 KiB
Go
package mctx
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"text/tabwriter"
|
|
)
|
|
|
|
// MaxStackSize indicates the maximum number of stack frames which will be
|
|
// stored when embedding stack traces in errors.
|
|
var MaxStackSize = 50
|
|
|
|
type ctxStackKey int
|
|
|
|
// Stacktrace represents a stack trace at a particular point in execution.
|
|
type Stacktrace struct {
|
|
frames []uintptr
|
|
}
|
|
|
|
// Frame returns the first frame in the stack.
|
|
func (s Stacktrace) Frame() runtime.Frame {
|
|
if len(s.frames) == 0 {
|
|
panic("cannot call Frame on empty stack")
|
|
}
|
|
|
|
frame, _ := runtime.CallersFrames([]uintptr(s.frames)).Next()
|
|
return frame
|
|
}
|
|
|
|
// Frames returns all runtime.Frame instances for this stack.
|
|
func (s Stacktrace) Frames() []runtime.Frame {
|
|
if len(s.frames) == 0 {
|
|
return nil
|
|
}
|
|
|
|
out := make([]runtime.Frame, 0, len(s.frames))
|
|
frames := runtime.CallersFrames([]uintptr(s.frames))
|
|
for {
|
|
frame, more := frames.Next()
|
|
out = append(out, frame)
|
|
if !more {
|
|
break
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// String returns a string representing the top-most frame of the stack.
|
|
func (s Stacktrace) String() string {
|
|
if len(s.frames) == 0 {
|
|
return ""
|
|
}
|
|
frame := s.Frame()
|
|
file, dir := filepath.Base(frame.File), filepath.Dir(frame.File)
|
|
dir = filepath.Base(dir) // only want the first dirname, ie the pkg name
|
|
return fmt.Sprintf("%s/%s:%d", dir, file, frame.Line)
|
|
}
|
|
|
|
// FullString returns the full stack trace.
|
|
func (s Stacktrace) FullString() string {
|
|
sb := new(strings.Builder)
|
|
tw := tabwriter.NewWriter(sb, 0, 4, 4, ' ', 0)
|
|
for _, frame := range s.Frames() {
|
|
file := fmt.Sprintf("%s:%d", frame.File, frame.Line)
|
|
fmt.Fprintf(tw, "%s\t%s\n", file, frame.Function)
|
|
}
|
|
if err := tw.Flush(); err != nil {
|
|
panic(err)
|
|
}
|
|
return sb.String()
|
|
}
|
|
|
|
// WithStack returns a Context with the current stacktrace embedded in it (as a
|
|
// Stacktrace type). If skip is non-zero it will skip that many frames from the
|
|
// top of the stack. The frame containing the WithStack call itself is always
|
|
// excluded.
|
|
func WithStack(ctx context.Context, skip int) context.Context {
|
|
stackSlice := make([]uintptr, MaxStackSize)
|
|
// incr skip once for WithStack, and once for runtime.Callers
|
|
l := runtime.Callers(skip+2, stackSlice)
|
|
stack := Stacktrace{frames: stackSlice[:l]}
|
|
|
|
return context.WithValue(ctx, ctxStackKey(0), stack)
|
|
}
|
|
|
|
// Stack returns the Stacktrace instance which was embedded by WithStack, or false if
|
|
// none ever was.
|
|
func Stack(ctx context.Context) (Stacktrace, bool) {
|
|
stack, ok := ctx.Value(ctxStackKey(0)).(Stacktrace)
|
|
return stack, ok
|
|
}
|