Skip zap stack frames in stacktrace output

config_test.go

@@ -50,7 +50,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" +
-				`go.uber.org/zap.Stack`,
+				`testing.\w+`,
 		},
 	}
 

field_test.go

@@ -154,6 +154,6 @@ 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.Contains(t, f.String, "zap.TestStackField", "Expected stacktrace to contain caller.")
+	assert.Equal(t, takeStacktrace(), f.String, "Unexpected stack trace")
 	assertCanBeReused(t, f)
 }

logger_test.go

@@ -380,26 +380,6 @@ func TestLoggerAddCallerFail(t *testing.T) {
 	})
 }
 
-func TestLoggerAddStacktrace(t *testing.T) {
-	assertHasStack := func(t testing.TB, obs observer.LoggedEntry) {
-		assert.Contains(t, obs.Entry.Stack, "zap.TestLoggerAddStacktrace", "Expected to find test function in stacktrace.")
-	}
-	withLogger(t, DebugLevel, opts(AddStacktrace(InfoLevel)), func(logger *Logger, logs *observer.ObservedLogs) {
-		logger.Debug("")
-		assert.Empty(
-			t,
-			logs.AllUntimed()[0].Entry.Stack,
-			"Unexpected stacktrack at DebugLevel.",
-		)
-
-		logger.Info("")
-		assertHasStack(t, logs.AllUntimed()[1])
-
-		logger.Warn("")
-		assertHasStack(t, logs.AllUntimed()[2])
-	})
-}
-
 func TestLoggerReplaceCore(t *testing.T) {
 	replace := WrapCore(func(zapcore.Core) zapcore.Core {
 		return zapcore.NewNopCore()

stacktrace.go

@@ -28,6 +28,8 @@ import (
 	"go.uber.org/zap/internal/bufferpool"
 )
 
+const _zapPackage = "go.uber.org/zap"
+
 var (
 	_stacktraceIgnorePrefixes = []string{
 		"runtime.goexit",
@@ -38,6 +40,11 @@ var (
 			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 {
@@ -60,11 +67,18 @@ func takeStacktrace() string {
 	}
 
 	i := 0
+	skipZapFrames := true // skip all consecutive zap frames at the beginning.
 	frames := runtime.CallersFrames(programCounters.pcs[:numFrames])
 	for frame, more := frames.Next(); more; frame, more = frames.Next() {
 		if shouldIgnoreStacktraceFunction(frame.Function) {
 			continue
 		}
+		if skipZapFrames && isZapFrame(frame.Function) {
+			continue
+		} else {
+			skipZapFrames = false
+		}
+
 		if i != 0 {
 			buffer.AppendByte('\n')
 		}
@@ -80,6 +94,24 @@ 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
+}
+
 func shouldIgnoreStacktraceFunction(function string) bool {
 	for _, prefix := range _stacktraceIgnorePrefixes {
 		if strings.HasPrefix(function, prefix) {
@@ -96,3 +128,11 @@ type programCounters struct {
 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

@@ -0,0 +1,161 @@
+// Copyright (c) 2016, 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package zap_test
+
+import (
+	"bytes"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"testing"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// _zapPackages are packages that we search for in the logging output to match a
+// zap stack frame. It is different from _zapStacktracePrefixes which  is only
+// intended to match on the function name, while this is on the full output
+// which includes filenames.
+var _zapPackages = []string{
+	"go.uber.org/zap.",
+	"go.uber.org/zap/zapcore.",
+}
+
+func TestStacktraceFiltersZapLog(t *testing.T) {
+	withLogger(t, func(logger *zap.Logger, out *bytes.Buffer) {
+		logger.Error("test log")
+		logger.Sugar().Error("sugar test log")
+
+		require.Contains(t, out.String(), "TestStacktraceFiltersZapLog", "Should not strip out non-zap import")
+		verifyNoZap(t, out.String())
+	})
+}
+
+func TestStacktraceFiltersZapMarshal(t *testing.T) {
+	withLogger(t, func(logger *zap.Logger, out *bytes.Buffer) {
+		marshal := func(enc zapcore.ObjectEncoder) error {
+			logger.Warn("marshal caused warn")
+			enc.AddString("f", "v")
+			return nil
+		}
+		logger.Error("test log", zap.Object("obj", zapcore.ObjectMarshalerFunc(marshal)))
+
+		logs := out.String()
+
+		// The marshal function (which will be under the test function) should not be stripped.
+		const marshalFnPrefix = "TestStacktraceFiltersZapMarshal."
+		require.Contains(t, logs, marshalFnPrefix, "Should not strip out marshal call")
+
+		// There should be no zap stack traces before that point.
+		marshalIndex := strings.Index(logs, marshalFnPrefix)
+		verifyNoZap(t, logs[:marshalIndex])
+
+		// After that point, there should be zap stack traces - we don't want to strip out
+		// the Marshal caller information.
+		for _, fnPrefix := range _zapPackages {
+			require.Contains(t, logs[marshalIndex:], fnPrefix, "Missing zap caller stack for Marshal")
+		}
+	})
+}
+
+func TestStacktraceFiltersVendorZap(t *testing.T) {
+	// We need to simulate a zap as a vendor library, so we're going to create a fake GOPATH
+	// and run the above test which will contain zap in the vendor directory.
+	withGoPath(t, func(goPath string) {
+		curDir, err := os.Getwd()
+		require.NoError(t, err, "Failed to get current directory")
+
+		testDir := filepath.Join(goPath, "src/go.uber.org/zap_test/")
+		vendorDir := filepath.Join(testDir, "vendor")
+		require.NoError(t, os.MkdirAll(testDir, 0777), "Failed to create source director")
+
+		curFile := getSelfFilename(t)
+		//copyFile(t, curFile, filepath.Join(testDir, curFile))
+		setupSymlink(t, curFile, filepath.Join(testDir, curFile))
+
+		// Set up symlinks for zap, and for any test dependencies.
+		setupSymlink(t, curDir, filepath.Join(vendorDir, "go.uber.org/zap"))
+		for _, testDep := range []string{"github.com/stretchr/testify"} {
+			setupSymlink(t, filepath.Join(curDir, "vendor", testDep), filepath.Join(vendorDir, testDep))
+		}
+
+		// Now run the above test which ensures we filter out zap stacktraces, but this time
+		// zap is in a vendor
+		cmd := exec.Command("go", "test", "-v", "-run", "TestStacktraceFiltersZap")
+		cmd.Dir = testDir
+		out, err := cmd.CombinedOutput()
+		require.NoError(t, err, "Failed to run test in vendor directory, output: %s", out)
+		assert.Contains(t, string(out), "PASS")
+	})
+}
+
+// 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)) {
+	buf := &bytes.Buffer{}
+	encoder := zapcore.NewConsoleEncoder(zap.NewDevelopmentEncoderConfig())
+	core := zapcore.NewCore(encoder, zapcore.AddSync(buf), zapcore.DebugLevel)
+	logger := zap.New(core, zap.AddStacktrace(zap.DebugLevel))
+	fn(logger, buf)
+}
+
+func verifyNoZap(t *testing.T, logs string) {
+	for _, fnPrefix := range _zapPackages {
+		require.NotContains(t, logs, fnPrefix, "Should not strip out marshal call")
+	}
+}
+
+func withGoPath(t *testing.T, f func(goPath string)) {
+	goPath, err := ioutil.TempDir("", "gopath")
+	require.NoError(t, err, "Failed to create temporary directory for GOPATH")
+	//defer os.RemoveAll(goPath)
+
+	os.Setenv("GOPATH", goPath)
+	defer os.Setenv("GOPATH", os.Getenv("GOPATH"))
+
+	f(goPath)
+}
+
+func getSelfFilename(t *testing.T) string {
+	_, file, _, ok := runtime.Caller(0)
+	require.True(t, ok, "Failed to get caller information to identify local file")
+
+	return filepath.Base(file)
+}
+
+func setupSymlink(t *testing.T, src, dst string) {
+	// Make sure the destination directory exists.
+	os.MkdirAll(filepath.Dir(dst), 0777)
+
+	// Get absolute path of the source for the symlink, otherwise we can create a symlink
+	// that uses relative paths.
+	srcAbs, err := filepath.Abs(src)
+	require.NoError(t, err, "Failed to get absolute path")
+
+	require.NoError(t, os.Symlink(srcAbs, dst), "Failed to set up symlink")
+}

stacktrace_test.go

@@ -35,7 +35,41 @@ func TestTakeStacktrace(t *testing.T) {
 	assert.Contains(
 		t,
 		lines[0],
-		"TestTakeStacktrace",
-		"Expected stacktrace to start with this test function, but top frame is %s.", lines[0],
+		"testing.",
+		"Expected stacktrace to start with the test runner (zap frames are filtered out) %s.", lines[0],
 	)
 }
+
+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",
+	}
+
+	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 BenchmarkTakeStacktrace(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		takeStacktrace()
+	}
+}