From 3786fbfa73026bbd445798c02f8b4195cff9d14f Mon Sep 17 00:00:00 2001
From: Olivier Poitrey <rs@rhapsodyk.net>
Date: Wed, 7 Feb 2018 21:55:53 -0800
Subject: [PATCH] Add support for stack trace extration of error fields

---
 context.go                   | 14 ++++++++
 event.go                     | 22 ++++++++++++
 globals.go                   |  7 ++++
 pkgerrors/stacktrace.go      | 66 ++++++++++++++++++++++++++++++++++++
 pkgerrors/stacktrace_test.go | 26 ++++++++++++++
 5 files changed, 135 insertions(+)
 create mode 100644 pkgerrors/stacktrace.go
 create mode 100644 pkgerrors/stacktrace_test.go

diff --git a/context.go b/context.go
index 4137287..b52bb0e 100644
--- a/context.go
+++ b/context.go
@@ -316,3 +316,17 @@ func (c Context) Caller() Context {
 	c.l = c.l.Hook(ch)
 	return c
 }
+
+type stackTraceHook struct{}
+
+func (sh stackTraceHook) Run(e *Event, level Level, msg string) {
+	e.Stack()
+}
+
+var sh = callerHook{}
+
+// Stack enables stack trace printing for the error passed to Err().
+func (c Context) Stack() Context {
+	c.l = c.l.Hook(sh)
+	return c
+}
diff --git a/event.go b/event.go
index 66e0489..05c4af9 100644
--- a/event.go
+++ b/event.go
@@ -27,6 +27,7 @@ type Event struct {
 	w     LevelWriter
 	level Level
 	done  func(msg string)
+	stack bool   // enable error stack trace
 	ch    []Hook // hooks from context
 	h     []Hook
 }
@@ -242,17 +243,38 @@ func (e *Event) Errs(key string, errs []error) *Event {
 
 // Err adds the field "error" with err as a string to the *Event context.
 // If err is nil, no field is added.
+//
 // To customize the key name, change zerolog.ErrorFieldName.
+//
+// If Stack() has been called before and zerolog.ErrorStackMarshaler is defined,
+// the err is passed to ErrorStackMarshaler and the result is appended to the
+// zerolog.ErrorStackFieldName.
 func (e *Event) Err(err error) *Event {
 	if e == nil {
 		return e
 	}
+	if e.stack && ErrorStackMarshaler != nil {
+		s := ErrorStackMarshaler(err)
+		if len(s) > 0 {
+			e.buf = append(json.AppendKey(e.buf, ErrorStackFieldName), s...)
+		}
+	}
 	if err != nil {
 		e.buf = json.AppendError(json.AppendKey(e.buf, ErrorFieldName), err)
 	}
 	return e
 }
 
+// Stack enables stack trace printing for the error passed to Err().
+//
+// ErrorStackMarshaler must be set for this method to do something.
+func (e *Event) Stack() *Event {
+	if e != nil {
+		e.stack = true
+	}
+	return e
+}
+
 // Bool adds the field key with val as a bool to the *Event context.
 func (e *Event) Bool(key string, b bool) *Event {
 	if e == nil {
diff --git a/globals.go b/globals.go
index b65d0cb..97eed50 100644
--- a/globals.go
+++ b/globals.go
@@ -19,6 +19,13 @@ var (
 	// CallerFieldName is the field name used for caller field.
 	CallerFieldName = "caller"
 
+	// ErrorStackFieldName is the field name used for error stacks.
+	ErrorStackFieldName = "stack"
+
+	// ErrorStackMarshaler extract the stack from err if any, and returns it as
+	// a marshaled JSON.
+	ErrorStackMarshaler func(err error) []byte
+
 	// TimeFieldFormat defines the time format of the Time field type.
 	// If set to an empty string, the time is formatted as an UNIX timestamp
 	// as integer.
diff --git a/pkgerrors/stacktrace.go b/pkgerrors/stacktrace.go
new file mode 100644
index 0000000..fff1c47
--- /dev/null
+++ b/pkgerrors/stacktrace.go
@@ -0,0 +1,66 @@
+package pkgerrors
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/internal/json"
+)
+
+var (
+	StackSourceFileName     = "source"
+	StackSourceLineName     = "line"
+	StackSourceFunctionName = "func"
+)
+
+// MarshalStack implements pkg/errors stack trace marshaling.
+//
+//   zerolog.ErrorStackMarshaler = MarshalStack
+func MarshalStack(err error) []byte {
+	type stackTracer interface {
+		StackTrace() errors.StackTrace
+	}
+	var st errors.StackTrace
+	if err, ok := err.(stackTracer); ok {
+		st = err.StackTrace()
+	} else {
+		return nil
+	}
+	return appendJSONStack(make([]byte, 0, 500), st)
+}
+
+func appendJSONStack(dst []byte, st errors.StackTrace) []byte {
+	buf := bytes.NewBuffer(make([]byte, 0, 100))
+	dst = append(dst, '[')
+	for i, frame := range st {
+		if i > 0 {
+			dst = append(dst, ',')
+		}
+
+		dst = append(dst, '{')
+
+		fmt.Fprintf(buf, "%s", frame)
+		dst = json.AppendString(dst, StackSourceFileName)
+		dst = append(dst, ':')
+		dst = json.AppendBytes(dst, buf.Bytes())
+		dst = append(dst, ',')
+		buf.Reset()
+
+		fmt.Fprintf(buf, "%d", frame)
+		dst = json.AppendString(dst, StackSourceLineName)
+		dst = append(dst, ':')
+		dst = json.AppendBytes(dst, buf.Bytes())
+		dst = append(dst, ',')
+		buf.Reset()
+
+		fmt.Fprintf(buf, "%n", frame)
+		dst = json.AppendString(dst, StackSourceFunctionName)
+		dst = append(dst, ':')
+		dst = json.AppendBytes(dst, buf.Bytes())
+
+		dst = append(dst, '}')
+	}
+	dst = append(dst, ']')
+	return dst
+}
diff --git a/pkgerrors/stacktrace_test.go b/pkgerrors/stacktrace_test.go
new file mode 100644
index 0000000..bee1df0
--- /dev/null
+++ b/pkgerrors/stacktrace_test.go
@@ -0,0 +1,26 @@
+package pkgerrors
+
+import (
+	"bytes"
+	"regexp"
+	"testing"
+
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog"
+)
+
+func TestLogStack(t *testing.T) {
+	zerolog.ErrorStackMarshaler = MarshalStack
+
+	out := &bytes.Buffer{}
+	log := zerolog.New(out)
+
+	err := errors.Wrap(errors.New("error message"), "from error")
+	log.Log().Stack().Err(err).Msg("")
+
+	got := out.String()
+	want := `\{"stack":\[\{"source":"stacktrace_test.go","line":"18","func":"TestLogStack"\},.*\],"error":"from error: error message"\}\n`
+	if ok, _ := regexp.MatchString(want, got); !ok {
+		t.Errorf("invalid log output:\ngot:  %v\nwant: %v", got, want)
+	}
+}