Share callstack b/w Caller and Stack

logger.go

@@ -296,15 +296,44 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry {
 
  // Thread the error output through to the CheckedEntry.
  ce.ErrorOutput = log.errorOutput
+
+ addStack := log.addStack.Enabled(ce.Entry.Level)
+ if !log.addCaller && !addStack {
+ return ce
+ }
+
+ // If a stack trace was requested, capture the entire stack.
+ stackLimit := 0
+
+ // For historical reasons, stack traces reported by Zap don't respect
+ // CallerSkip. See also #727.
+ stackSkip := callerSkipOffset
+
+ if !addStack {
+ stackLimit = 1
+ stackSkip += log.callerSkip
+ }
+
+ programCounters := newProgramCounters()
+ defer programCounters.Release()
+
+ pcs := programCounters.Callers(stackSkip, stackLimit)
  if log.addCaller {
- ce.Entry.Caller = zapcore.NewEntryCaller(runtime.Caller(log.callerSkip + callerSkipOffset))
+ if len(pcs) == 0 {
+ ce.Entry.Caller = zapcore.NewEntryCaller(0, "", 0, false)
+ } else {
+ frame, _ := runtime.CallersFrames(pcs).Next()
+ ce.Entry.Caller = zapcore.NewEntryCaller(frame.PC, frame.File, frame.Line, true)
+ }
+
  if !ce.Entry.Caller.Defined {
  fmt.Fprintf(log.errorOutput, "%v Logger.check error: failed to get caller\n", time.Now().UTC())
  log.errorOutput.Sync()
  }
  }
- if log.addStack.Enabled(ce.Entry.Level) {
- ce.Entry.Stack = Stack("").String
+
+ if addStack {
+ ce.Entry.Stack = makeStacktrace(pcs, true /* skip zap */)
  }
 
  return ce

logger_bench_test.go

@@ -179,6 +179,41 @@ func BenchmarkAddCallerHook(b *testing.B) {
  })
 }
 
+func BenchmarkAddStack(b *testing.B) {
+ logger := New(
+ zapcore.NewCore(
+ zapcore.NewJSONEncoder(NewProductionConfig().EncoderConfig),
+ &ztest.Discarder{},
+ InfoLevel,
+ ),
+ AddStacktrace(InfoLevel),
+ )
+ b.ResetTimer()
+ b.RunParallel(func(pb *testing.PB) {
+ for pb.Next() {
+ logger.Info("foo")
+ }
+ })
+}
+
+func BenchmarkAddCallerAndStack(b *testing.B) {
+ logger := New(
+ zapcore.NewCore(
+ zapcore.NewJSONEncoder(NewProductionConfig().EncoderConfig),
+ &ztest.Discarder{},
+ InfoLevel,
+ ),
+ AddCaller(),
+ AddStacktrace(InfoLevel),
+ )
+ b.ResetTimer()
+ b.RunParallel(func(pb *testing.PB) {
+ for pb.Next() {
+ logger.Info("foo")
+ }
+ })
+}
+
 func Benchmark10Fields(b *testing.B) {
  withBenchedLogger(b, func(log *Logger) {
  log.Info("Ten fields, passed at the log site.",

stacktrace.go

@@ -31,45 +31,34 @@ import (
 const _zapPackage = "go.uber.org/zap"
 
 var (
- _stacktracePool = sync.Pool{
- New: func() interface{} {
- return newProgramCounters(64)
- },
- }
-
  // We add "." and "/" suffixes to the package name to ensure we only match
  // the exact package and not any package with the same prefix.
  _zapStacktracePrefixes = addPrefix(_zapPackage, ".", "/")
  _zapStacktraceVendorContains = addPrefix("/vendor/", _zapStacktracePrefixes...)
 )
 
 func takeStacktrace() string {
+ pcs := newProgramCounters()
+ defer pcs.Release()
+
+ return makeStacktrace(
+ pcs.Callers(1 /* skip */, 0 /* limit */),
+ true, // skip zap
+ )
+}
+
+func makeStacktrace(pcs []uintptr, skipZapFrames bool) string {
  buffer := bufferpool.Get()
  defer buffer.Free()
- programCounters := _stacktracePool.Get().(*programCounters)
- defer _stacktracePool.Put(programCounters)
-
- var numFrames int
- for {
- // Skip the call to runtime.Counters and takeStacktrace so that the
- // program counters start at the caller of takeStacktrace.
- numFrames = runtime.Callers(2, programCounters.pcs)
- if numFrames < len(programCounters.pcs) {
- break
- }
- // Don't put the too-short counter slice back into the pool; this lets
- // the pool adjust if we consistently take deep stacktraces.
- programCounters = newProgramCounters(len(programCounters.pcs) * 2)
- }
 
  i := 0
- skipZapFrames := true // skip all consecutive zap frames at the beginning.
- frames := runtime.CallersFrames(programCounters.pcs[:numFrames])
+ frames := runtime.CallersFrames(pcs)
 
  // Note: On the last iteration, frames.Next() returns false, with a valid
  // frame, but we ignore this frame. The last frame is a a runtime frame which
  // adds noise, since it's only either runtime.main or runtime.goexit.
  for frame, more := frames.Next(); more; frame, more = frames.Next() {
+ // skip all consecutive zap frames at the beginning.
  if skipZapFrames && isZapFrame(frame.Function) {
  continue
  } else {
@@ -113,8 +102,55 @@ type programCounters struct {
  pcs []uintptr
 }
 
-func newProgramCounters(size int) *programCounters {
- return &programCounters{make([]uintptr, size)}
+var _stacktracePool = sync.Pool{
+ New: func() interface{} {
+ return &programCounters{make([]uintptr, 64)}
+ },
+}
+
+func newProgramCounters() *programCounters {
+ return _stacktracePool.Get().(*programCounters)
+}
+
+func (pcs *programCounters) Release() {
+ _stacktracePool.Put(pcs)
+}
+
+// Callers captures program counters of callers of the current function.
+//
+// skip specifies the number of functions to skip.
+// limit specifies how deep into the stack we go, zero meaning unlimited.
+func (pcs *programCounters) Callers(skip int, limit int) []uintptr {
+ pc := pcs.pcs
+ if limit > 0 {
+ // The shortened slice will never be placed back into the
+ // pool, but the longer one will.
+ if limit < len(pc) {
+ pc = pc[:limit]
+ } else {
+ pc = make([]uintptr, limit)
+ pcs.pcs = pc
+ }
+ }
+
+ var numFrames int
+ for {
+ // Skip the call to runtime.Callers and programCounter.Callers
+ // so that the program counters start at our caller.
+ numFrames = runtime.Callers(skip+2, pc)
+
+ // All required frames have been captured if the limit was
+ // non-zero or if the stack depth is smaller than the size of
+ // the slice.
+ if limit > 0 || numFrames < len(pc) {
+ return pc[:numFrames]
+ }
+
+ // Don't put the too-short counter slice back into the pool; this lets
+ // the pool adjust if we consistently take deep stacktraces.
+ pc = make([]uintptr, len(pcs.pcs)*2)
+ pcs.pcs = pc
+ }
 }
 
 func addPrefix(prefix string, ss ...string) []string {

stacktrace_test.go

@@ -73,3 +73,56 @@ func BenchmarkTakeStacktrace(b *testing.B) {
  takeStacktrace()
  }
 }
+
+func TestProgramCountersCallersDeepStack(t *testing.T) {
+ tests := []struct {
+ desc string
+ stackDepth int
+ limit int
+ wantOccurrences int
+ }{
+ {
+ desc: "large limit",
+ stackDepth: 128,
+ limit: 128,
+ wantOccurrences: 127,
+ // one of the slots is taken by the closure we pass to
+ // withStackDepth.
+ },
+ {
+ desc: "unlimited",
+ stackDepth: 500,
+ wantOccurrences: 500,
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.desc, func(t *testing.T) {
+ const withStackDepthName = "go.uber.org/zap.withStackDepth"
+
+ var trace string
+ withStackDepth(tt.stackDepth, func() {
+ pcs := newProgramCounters().Callers(1, tt.limit)
+ trace = makeStacktrace(pcs, false /* don't skip zap */)
+ })
+
+ for found := 0; found < tt.wantOccurrences; found++ {
+ i := strings.Index(trace, withStackDepthName)
+ if i < 0 {
+ t.Fatalf(`expected %d occurrences of %q, found %d`, tt.wantOccurrences, withStackDepthName, found)
+ }
+
+ trace = trace[i+len(withStackDepthName):]
+ }
+ })
+ }
+
+}
+
+func withStackDepth(depth int, f func()) {
+ if depth > 0 {
+ withStackDepth(depth-1, f)
+ } else {
+ f()
+ }
+}