Honor CallerSkip when taking stack traces & add the StackSkip field (u…

…ber-go#843)

* Honor `CallerSkip` when taking a stack trace. Both the caller and stack trace will now point to the same frame.
* Add `StackSkip` which is similar to `Stack` but allows skipping frames from the top of the stack.

This removes the internal behavior of skipping zap specific stack frames when taking a stack trace, and instead relies on the same behavior used to calculate the number of frames to skip for getting the caller to also skip frames in the stack trace.

Fixes uber-go#512, fixes uber-go#727

config_test.go

@@ -52,7 +52,7 @@ func TestConfig(t *testing.T) {
 			expectRe: "DEBUG\tzap/config_test.go:" + `\d+` + "\tdebug\t" + `{"k": "v", "z": "zz"}` + "\n" +
 				"INFO\tzap/config_test.go:" + `\d+` + "\tinfo\t" + `{"k": "v", "z": "zz"}` + "\n" +
 				"WARN\tzap/config_test.go:" + `\d+` + "\twarn\t" + `{"k": "v", "z": "zz"}` + "\n" +
-				`testing.\w+`,
+				`go.uber.org/zap.TestConfig.\w+`,
 		},
 	}
 

field.go

@@ -364,11 +364,17 @@ func Timep(key string, val *time.Time) Field {
 // expensive (relatively speaking); this function both makes an allocation and
 // takes about two microseconds.
 func Stack(key string) Field {
+	return StackSkip(key, 1) // skip Stack
+}
+
+// StackSkip constructs a field similarly to Stack, but also skips the given
+// number of frames from the top of the stacktrace.
+func StackSkip(key string, skip int) Field {
 	// Returning the stacktrace as a string costs an allocation, but saves us
 	// from expanding the zapcore.Field union struct to include a byte slice. Since
 	// taking a stacktrace is already so expensive (~10us), the extra allocation
 	// is okay.
-	return String(key, takeStacktrace())
+	return String(key, takeStacktrace(skip+1)) // skip StackSkip
 }
 
 // Duration constructs a field with the given key and value. The encoder

field_test.go

@@ -23,6 +23,7 @@ package zap
 import (
 	"math"
 	"net"
+	"regexp"
 	"sync"
 	"testing"
 	"time"
@@ -259,6 +260,24 @@ func TestStackField(t *testing.T) {
 	f := Stack("stacktrace")
 	assert.Equal(t, "stacktrace", f.Key, "Unexpected field key.")
 	assert.Equal(t, zapcore.StringType, f.Type, "Unexpected field type.")
-	assert.Equal(t, takeStacktrace(), f.String, "Unexpected stack trace")
+	r := regexp.MustCompile(`field_test.go:(\d+)`)
+	assert.Equal(t, r.ReplaceAllString(takeStacktrace(0), "field_test.go"), r.ReplaceAllString(f.String, "field_test.go"), "Unexpected stack trace")
+	assertCanBeReused(t, f)
+}
+
+func TestStackSkipField(t *testing.T) {
+	f := StackSkip("stacktrace", 0)
+	assert.Equal(t, "stacktrace", f.Key, "Unexpected field key.")
+	assert.Equal(t, zapcore.StringType, f.Type, "Unexpected field type.")
+	r := regexp.MustCompile(`field_test.go:(\d+)`)
+	assert.Equal(t, r.ReplaceAllString(takeStacktrace(0), "field_test.go"), r.ReplaceAllString(f.String, "field_test.go"), f.String, "Unexpected stack trace")
+	assertCanBeReused(t, f)
+}
+
+func TestStackSkipFieldWithSkip(t *testing.T) {
+	f := StackSkip("stacktrace", 1)
+	assert.Equal(t, "stacktrace", f.Key, "Unexpected field key.")
+	assert.Equal(t, zapcore.StringType, f.Type, "Unexpected field type.")
+	assert.Equal(t, takeStacktrace(1), f.String, "Unexpected stack trace")
 	assertCanBeReused(t, f)
 }

logger.go

@@ -312,7 +312,7 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry {
 		}
 	}
 	if log.addStack.Enabled(ce.Entry.Level) {
-		ce.Entry.Stack = Stack("").String
+		ce.Entry.Stack = StackSkip("", log.callerSkip+callerSkipOffset).String
 	}
 
 	return ce

stacktrace.go

@@ -22,38 +22,30 @@ package zap
 
 import (
 	"runtime"
-	"strings"
 	"sync"
 
 	"go.uber.org/zap/internal/bufferpool"
 )
 
-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 {
+func takeStacktrace(skip int) 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
+		// Skip the call to runtime.Callers and takeStacktrace so that the
 		// program counters start at the caller of takeStacktrace.
-		numFrames = runtime.Callers(2, programCounters.pcs)
+		numFrames = runtime.Callers(skip+2, programCounters.pcs)
 		if numFrames < len(programCounters.pcs) {
 			break
 		}
@@ -63,19 +55,12 @@ func takeStacktrace() string {
 	}
 
 	i := 0
-	skipZapFrames := true // skip all consecutive zap frames at the beginning.
 	frames := runtime.CallersFrames(programCounters.pcs[:numFrames])
 
 	// 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() {
-		if skipZapFrames && isZapFrame(frame.Function) {
-			continue
-		} else {
-			skipZapFrames = false
-		}
-
 		if i != 0 {
 			buffer.AppendByte('\n')
 		}
@@ -91,36 +76,10 @@ func takeStacktrace() string {
 	return buffer.String()
 }
 
-func isZapFrame(function string) bool {
-	for _, prefix := range _zapStacktracePrefixes {
-		if strings.HasPrefix(function, prefix) {
-			return true
-		}
-	}
-
-	// We can't use a prefix match here since the location of the vendor
-	// directory affects the prefix. Instead we do a contains match.
-	for _, contains := range _zapStacktraceVendorContains {
-		if strings.Contains(function, contains) {
-			return true
-		}
-	}
-
-	return false
-}
-
 type programCounters struct {
 	pcs []uintptr
 }
 
 func newProgramCounters(size int) *programCounters {
 	return &programCounters{make([]uintptr, size)}
 }
-
-func addPrefix(prefix string, ss ...string) []string {
-	withPrefix := make([]string, len(ss))
-	for i, s := range ss {
-		withPrefix[i] = prefix + s
-	}
-	return withPrefix
-}

stacktrace_ext_test.go

@@ -120,6 +120,30 @@ func TestStacktraceFiltersVendorZap(t *testing.T) {
 	})
 }
 
+func TestStacktraceWithoutCallerSkip(t *testing.T) {
+	withLogger(t, func(logger *zap.Logger, out *bytes.Buffer) {
+		func() {
+			logger.Error("test log")
+		}()
+
+		require.Contains(t, out.String(), "TestStacktraceWithoutCallerSkip.", "Should not skip too much")
+		verifyNoZap(t, out.String())
+	})
+}
+
+func TestStacktraceWithCallerSkip(t *testing.T) {
+	withLogger(t, func(logger *zap.Logger, out *bytes.Buffer) {
+		logger = logger.WithOptions(zap.AddCallerSkip(2))
+		func() {
+			logger.Error("test log")
+		}()
+
+		require.NotContains(t, out.String(), "TestStacktraceWithCallerSkip.", "Should skip as requested by caller skip")
+		require.Contains(t, out.String(), "TestStacktraceWithCallerSkip", "Should not skip too much")
+		verifyNoZap(t, out.String())
+	})
+}
+
 // withLogger sets up a logger with a real encoder set up, so that any marshal functions are called.
 // The inbuilt observer does not call Marshal for objects/arrays, which we need for some tests.
 func withLogger(t *testing.T, fn func(logger *zap.Logger, out *bytes.Buffer)) {

stacktrace_test.go

@@ -29,47 +29,46 @@ import (
 )
 
 func TestTakeStacktrace(t *testing.T) {
-	trace := takeStacktrace()
+	trace := takeStacktrace(0)
 	lines := strings.Split(trace, "\n")
-	require.True(t, len(lines) > 0, "Expected stacktrace to have at least one frame.")
+	require.NotEmpty(t, lines, "Expected stacktrace to have at least one frame.")
 	assert.Contains(
 		t,
 		lines[0],
-		"testing.",
-		"Expected stacktrace to start with the test runner (zap frames are filtered out) %s.", lines[0],
+		"go.uber.org/zap.TestTakeStacktrace",
+		"Expected stacktrace to start with the test.",
 	)
 }
 
-func TestIsZapFrame(t *testing.T) {
-	zapFrames := []string{
-		"go.uber.org/zap.Stack",
-		"go.uber.org/zap.(*SugaredLogger).log",
-		"go.uber.org/zap/zapcore.(ArrayMarshalerFunc).MarshalLogArray",
-		"github.com/uber/tchannel-go/vendor/go.uber.org/zap.Stack",
-		"github.com/uber/tchannel-go/vendor/go.uber.org/zap.(*SugaredLogger).log",
-		"github.com/uber/tchannel-go/vendor/go.uber.org/zap/zapcore.(ArrayMarshalerFunc).MarshalLogArray",
-	}
-	nonZapFrames := []string{
-		"github.com/uber/tchannel-go.NewChannel",
-		"go.uber.org/not-zap.New",
-		"go.uber.org/zapext.ctx",
-		"go.uber.org/zap_ext/ctx.New",
-	}
+func TestTakeStacktraceWithSkip(t *testing.T) {
+	trace := takeStacktrace(1)
+	lines := strings.Split(trace, "\n")
+	require.NotEmpty(t, lines, "Expected stacktrace to have at least one frame.")
+	assert.Contains(
+		t,
+		lines[0],
+		"testing.",
+		"Expected stacktrace to start with the test runner (skipping our own frame).",
+	)
+}
 
-	t.Run("zap frames", func(t *testing.T) {
-		for _, f := range zapFrames {
-			require.True(t, isZapFrame(f), f)
-		}
-	})
-	t.Run("non-zap frames", func(t *testing.T) {
-		for _, f := range nonZapFrames {
-			require.False(t, isZapFrame(f), f)
-		}
-	})
+func TestTakeStacktraceWithSkipInnerFunc(t *testing.T) {
+	var trace string
+	func() {
+		trace = takeStacktrace(2)
+	}()
+	lines := strings.Split(trace, "\n")
+	require.NotEmpty(t, lines, "Expected stacktrace to have at least one frame.")
+	assert.Contains(
+		t,
+		lines[0],
+		"testing.",
+		"Expected stacktrace to start with the test function (skipping the test function).",
+	)
 }
 
 func BenchmarkTakeStacktrace(b *testing.B) {
 	for i := 0; i < b.N; i++ {
-		takeStacktrace()
+		takeStacktrace(0)
 	}
 }