stack.go

package fault

import (
	"errors"
	"fmt"
	"runtime"
	"runtime/debug"
	"strings"
	"sync"
)

const maxStackDepth = 32

// Stack is an abstraction for a call stack captured with CaptureStack,
// typically when handling a runtime error.
//
// This implementation shifts costly processing from capture time to reporting
// time because the latter may never happen.
type Stack struct {
	pc []uintptr
}

// CaptureStack captures the current call stack.
//
// The skip argument specifies how many innermost frames to ignore; skip = 0
// omits CaptureStack but includes its caller.
func CaptureStack(skip int) Stack {
	var buf [maxStackDepth]uintptr
	n := runtime.Callers(skip+2, buf[:])
	return Stack{pc: buf[:n]}
}

// Frames returns the stack frames in outward order
func (s Stack) Frames() []runtime.Frame {
	var res []runtime.Frame
	frames := runtime.CallersFrames(s.pc)
	for {
		frame, more := frames.Next()
		if frame != (runtime.Frame{}) {
			res = append(res, frame)
		}
		if !more {
			break
		}
	}
	return res
}

// ShortenFunction is a convenience function that takes a fully qualified
// function name as found in runtime.Frame.Function and removes the path to the
// main module of the running binary. For example, in a binary built from
// module "example.com/app", the name "example.com/app/service.Run" becomes
// "service.Run". Any other function name is returned unchanged.
func ShortenFunction(fn string) string {
	return strings.TrimPrefix(fn, pathPrefix())
}

var pathPrefix = sync.OnceValue(func() string {
	info, ok := debug.ReadBuildInfo()
	if !ok {
		return ""
	}
	return info.Main.Path + "/"
})

// String formats the stack as text similar to the way the standard library does
// it. File names are as recorded at build time: absolute paths on the build
// machine, or module-relative ones if the binary was built with -trimpath.
func (s Stack) String() string {
	var entries []string
	for _, frame := range s.Frames() {
		entries = append(entries, fmt.Sprintf("%s\n\t%s:%d", ShortenFunction(frame.Function), frame.File, frame.Line))
	}
	return strings.Join(entries, "\n")
}

// WithStack is any error with an associated Stack
type WithStack interface {
	error
	Stack() Stack
}

type stackError struct {
	wrapped error
	stack   *Stack // pointer to keep the struct comparable
}

func (se stackError) Error() string {
	return se.wrapped.Error()
}

func (se stackError) Unwrap() error {
	return se.wrapped
}

func (se stackError) Stack() Stack {
	return *se.stack
}

// WrapStack wraps the error with a Stack unless it already has one. The skip
// argument specifies how many innermost frames to omit; skip = 0 omits WrapStack
// but includes its caller.
func WrapStack(err error, skip int) error {
	if err == nil {
		return nil
	}
	if _, ok := errors.AsType[WithStack](err); ok {
		return err
	}
	stack := CaptureStack(skip + 1)
	return stackError{wrapped: err, stack: &stack}
}