Add skip level support to the AddCaller hook; Add AddCallerWithSkip a…

…nd AddCallersWithSkip hooks

field.go

@@ -22,6 +22,7 @@ package zap
 
 import (
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"math"
 	"runtime"
@@ -48,6 +49,8 @@ const (
 	skipType
 )
 
+var errCaller = errors.New("failed to get caller")
+
 // A Field is a marshaling operation used to add a key-value pair to a logger's
 // context. Most fields are lazily marshaled, so it's inexpensive to add fields to
 // disabled debug-level log statements.
@@ -205,16 +208,20 @@ func Callers(skip int) Field {
 //
 // The marshaled object has "pc", "file", and "line" fields, unless a caller
 // couldn't be resolved.
-func Caller(skip int) Field {
+func Caller(skip int) (Field, error) {
 	pc, file, line, ok := runtime.Caller(skip + 2)
+	if !ok {
+		return Skip(), errCaller
+	}
+
 	return Marshaler("caller", LogMarshalerFunc(func(kv KeyValue) error {
 		if ok {
 			kv.AddUintptr("pc", pc)
 			kv.AddString("file", file)
 			kv.AddInt("line", line)
 		}
 		return nil
-	}))
+	})), nil
 }
 
 // Duration constructs a Field with the given key and value. It represents

field_test.go

@@ -66,6 +66,15 @@ func assertNotEqualFieldJSON(t testing.TB, expected string, field Field) {
 		"Unexpected JSON output after applying field %+v.", field)
 }
 
+func assertFieldJSONRegEx(t testing.TB, expected string, field Field) {
+	enc := newJSONEncoder()
+	defer enc.Free()
+
+	field.AddTo(enc)
+	assert.Regexp(t, expected, string(enc.bytes),
+		"Unexpected JSON output after applying field %+v.", field)
+}
+
 func assertCanBeReused(t testing.TB, field Field) {
 	var wg sync.WaitGroup
 
@@ -216,6 +225,16 @@ func TestStackField(t *testing.T) {
 	assert.Contains(t, output[13:], "zap.TestStackField", "Expected stacktrace to contain caller.")
 }
 
+func TestCaller(t *testing.T) {
+	// Caller() makes an assupmtion that it is called within zap's hook mechanism.
+	// This is not true for this test, so pass -1 to log this function as the caller.
+	c, _ := Caller(-1)
+	assertFieldJSONRegEx(t, "\"caller\":{\"pc\":[0-9]+,\"file\":\".+field_test.go\",\"line\":[0-9]+}", c)
+
+	c2, _ := Caller(-1)
+	assertCanBeReused(t, c2)
+}
+
 func TestUnknownField(t *testing.T) {
 	enc := NewJSONEncoder()
 	defer enc.Free()

hook.go

@@ -20,19 +20,14 @@
 
 package zap
 
-import (
-	"errors"
-	"path/filepath"
-	"runtime"
-	"strconv"
-)
+import "errors"
 
 var (
 	errHookNilEntry = errors.New("can't call a hook on a nil *Entry")
-	errCaller       = errors.New("failed to get caller")
 	// Skip Caller, Logger.log, and the leveled Logger method when using
 	// runtime.Caller.
-	_callerSkip = 3
+	_callerSkip  = 2
+	_callersSkip = 3 // see golang docs for differences between runtime.Caller(s)
 )
 
 // A Hook is executed each time the logger writes an Entry. It can modify the
@@ -48,31 +43,45 @@ func (h Hook) apply(m *Meta) {
 	m.Hooks = append(m.Hooks, h)
 }
 
-// AddCaller configures the Logger to annotate each message with the filename
-// and line number of zap's caller.
+// AddCaller configures the Logger to annotate each log with
+// a "caller" field with the values as the filename,
+// line number and program counter.
 func AddCaller() Option {
+	return AddCallerWithSkip(_callerSkip)
+}
+
+// AddCallerWithSkip - see desc above
+// The callerSkip parameter helps specifies which stack frame above the
+// runtime.Caller() to capture.
+// Using _callerSkip will return the caller of the leveled logger method.
+func AddCallerWithSkip(callerSkip int) Option {
 	return Hook(func(e *Entry) error {
 		if e == nil {
 			return errHookNilEntry
 		}
-		_, filename, line, ok := runtime.Caller(_callerSkip)
-		if !ok {
-			return errCaller
+
+		var field Field
+		var err error
+		if field, err = Caller(callerSkip); err != nil {
+			return err
 		}
 
-		// Re-use a buffer from the pool.
-		enc := jsonPool.Get().(*jsonEncoder)
-		enc.truncate()
-		buf := enc.bytes
-		buf = append(buf, filepath.Base(filename)...)
-		buf = append(buf, ':')
-		buf = strconv.AppendInt(buf, int64(line), 10)
-		buf = append(buf, ':', ' ')
-		buf = append(buf, e.Message...)
+		field.AddTo(e.Fields())
+		return nil
+	})
+}
 
-		newMsg := string(buf)
-		enc.Free()
-		e.Message = newMsg
+// AddCallersWithSkip records the caller's call stack for all messages at or above a given level.
+// The callerskip option helps select the target caller.  Using a value of _callerSkip will select
+// the target caller as the caller of the leveled logger method.
+func AddCallersWithSkip(callerSkip int, lvl Level) Option {
+	return Hook(func(e *Entry) error {
+		if e == nil {
+			return errHookNilEntry
+		}
+		if e.Level >= lvl {
+			Callers(callerSkip).AddTo(e.Fields())
+		}
 		return nil
 	})
 }

hook_test.go

@@ -33,24 +33,47 @@ func TestHookAddCaller(t *testing.T) {
 	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), AddCaller())
 	logger.Info("Callers.")
 
-	re := regexp.MustCompile(`"msg":"hook_test.go:[\d]+: Callers\."`)
+	re := regexp.MustCompile(`"msg":"Callers\.","caller":{"pc":\d+,"file":".+hook_test.go","line":\d+}`)
+	assert.Regexp(t, re, buf.Stripped(), "Expected to find package name and file name in output.")
+}
+
+func TestHookAddCallerWithSkip(t *testing.T) {
+	buf := &testBuffer{}
+	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), AddCallerWithSkip(_callerSkip))
+	logger.Info("Callers.")
+
+	re := regexp.MustCompile(`"msg":"Callers\.","caller":{"pc":\d+,"file":".+hook_test.go","line":\d+}`)
 	assert.Regexp(t, re, buf.Stripped(), "Expected to find package name and file name in output.")
 }
 
 func TestHookAddCallerFail(t *testing.T) {
 	buf := &testBuffer{}
 	errBuf := &testBuffer{}
 
-	originalSkip := _callerSkip
-	_callerSkip = 1e3
-	defer func() { _callerSkip = originalSkip }()
-
-	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), ErrorOutput(errBuf), AddCaller())
+	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), ErrorOutput(errBuf), AddCallerWithSkip(1e3))
 	logger.Info("Failure.")
 	assert.Regexp(t, `hook error: failed to get caller`, errBuf.String(), "Didn't find expected failure message.")
 	assert.Contains(t, buf.String(), `"msg":"Failure."`, "Expected original message to survive failures in runtime.Caller.")
 }
 
+func TestHookAddCallersWithSkip(t *testing.T) {
+	buf := &testBuffer{}
+	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), AddCallersWithSkip(_callersSkip, InfoLevel))
+
+	logger.Info("Callers.")
+	output := buf.String()
+	require.Contains(t, output, "zap/hook_test.go", "Expected to find test file in callers.")
+	assert.Contains(t, output, `"callers":`, "Callers added under an unexpected key.")
+
+	buf.Reset()
+	logger.Warn("Callers.")
+	assert.Contains(t, buf.String(), `"callers":`, "Expected to include callers at Warn level.")
+
+	buf.Reset()
+	logger.Debug("No callers.")
+	assert.NotContains(t, buf.String(), "Unexpected stacktrace at Debug level.")
+}
+
 func TestHookAddStacks(t *testing.T) {
 	buf := &testBuffer{}
 	logger := New(NewJSONEncoder(), DebugLevel, Output(buf), AddStacks(InfoLevel))