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()
+	}
+}