From dc8ffe785c7fc9a74eeb5241814d77f1c5fb5e58 Mon Sep 17 00:00:00 2001
From: Chris Stockton <chrisstocktonaz@gmail.com>
Date: Tue, 11 Sep 2018 17:01:16 +0800
Subject: [PATCH] reduce allocations when printing stack traces (#149)

---
 bench_test.go | 49 ++++++++++++++++++++++++++++++++++++-
 stack.go      | 68 ++++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 102 insertions(+), 15 deletions(-)

diff --git a/bench_test.go b/bench_test.go
index 903b5f2..c906870 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -25,7 +25,7 @@ func yesErrors(at, depth int) error {
 
 // GlobalE is an exported global to store the result of benchmark results,
 // preventing the compiler from optimising the benchmark functions away.
-var GlobalE error
+var GlobalE interface{}
 
 func BenchmarkErrors(b *testing.B) {
 	type run struct {
@@ -61,3 +61,50 @@ func BenchmarkErrors(b *testing.B) {
 		})
 	}
 }
+
+func BenchmarkStackFormatting(b *testing.B) {
+	type run struct {
+		stack  int
+		format string
+	}
+	runs := []run{
+		{10, "%s"},
+		{10, "%v"},
+		{10, "%+v"},
+		{30, "%s"},
+		{30, "%v"},
+		{30, "%+v"},
+		{60, "%s"},
+		{60, "%v"},
+		{60, "%+v"},
+	}
+
+	var stackStr string
+	for _, r := range runs {
+		name := fmt.Sprintf("%s-stack-%d", r.format, r.stack)
+		b.Run(name, func(b *testing.B) {
+			err := yesErrors(0, r.stack)
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				stackStr = fmt.Sprintf(r.format, err)
+			}
+			b.StopTimer()
+		})
+	}
+
+	for _, r := range runs {
+		name := fmt.Sprintf("%s-stacktrace-%d", r.format, r.stack)
+		b.Run(name, func(b *testing.B) {
+			err := yesErrors(0, r.stack)
+			st := err.(*fundamental).stack.StackTrace()
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				stackStr = fmt.Sprintf(r.format, st)
+			}
+			b.StopTimer()
+		})
+	}
+	GlobalE = stackStr
+}
diff --git a/stack.go b/stack.go
index cbebb62..bb1e6a8 100644
--- a/stack.go
+++ b/stack.go
@@ -1,10 +1,12 @@
 package errors
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"path"
 	"runtime"
+	"strconv"
 	"strings"
 )
 
@@ -72,6 +74,11 @@ func (f Frame) line() int {
 //          GOPATH separated by \n\t (<funcname>\n\t<path>)
 //    %+v   equivalent to %+s:%d
 func (f Frame) Format(s fmt.State, verb rune) {
+	f.format(s, s, verb)
+}
+
+// format allows stack trace printing calls to be made with a bytes.Buffer.
+func (f Frame) format(w io.Writer, s fmt.State, verb rune) {
 	switch verb {
 	case 's':
 		switch {
@@ -79,23 +86,25 @@ func (f Frame) Format(s fmt.State, verb rune) {
 			pc := f.pc()
 			fn := runtime.FuncForPC(pc)
 			if fn == nil {
-				io.WriteString(s, "unknown")
+				io.WriteString(w, "unknown")
 			} else {
 				file, _ := fn.FileLine(pc)
-				fmt.Fprintf(s, "%s\n\t%s", fn.Name(), file)
+				io.WriteString(w, fn.Name())
+				io.WriteString(w, "\n\t")
+				io.WriteString(w, file)
 			}
 		default:
-			io.WriteString(s, path.Base(f.file()))
+			io.WriteString(w, path.Base(f.file()))
 		}
 	case 'd':
-		fmt.Fprintf(s, "%d", f.line())
+		io.WriteString(w, strconv.Itoa(f.line()))
 	case 'n':
 		name := runtime.FuncForPC(f.pc()).Name()
-		io.WriteString(s, funcname(name))
+		io.WriteString(w, funcname(name))
 	case 'v':
-		f.Format(s, 's')
-		io.WriteString(s, ":")
-		f.Format(s, 'd')
+		f.format(w, s, 's')
+		io.WriteString(w, ":")
+		f.format(w, s, 'd')
 	}
 }
 
@@ -111,23 +120,50 @@ type StackTrace []Frame
 //
 //    %+v   Prints filename, function, and line number for each Frame in the stack.
 func (st StackTrace) Format(s fmt.State, verb rune) {
+	var b bytes.Buffer
 	switch verb {
 	case 'v':
 		switch {
 		case s.Flag('+'):
-			for _, f := range st {
-				fmt.Fprintf(s, "\n%+v", f)
+			b.Grow(len(st) * stackMinLen)
+			for _, fr := range st {
+				b.WriteByte('\n')
+				fr.format(&b, s, verb)
 			}
 		case s.Flag('#'):
-			fmt.Fprintf(s, "%#v", []Frame(st))
+			fmt.Fprintf(&b, "%#v", []Frame(st))
 		default:
-			fmt.Fprintf(s, "%v", []Frame(st))
+			st.formatSlice(&b, s, verb)
 		}
 	case 's':
-		fmt.Fprintf(s, "%s", []Frame(st))
+		st.formatSlice(&b, s, verb)
+	}
+	io.Copy(s, &b)
+}
+
+// formatSlice will format this StackTrace into the given buffer as a slice of
+// Frame, only valid when called with '%s' or '%v'.
+func (st StackTrace) formatSlice(b *bytes.Buffer, s fmt.State, verb rune) {
+	b.WriteByte('[')
+	if len(st) == 0 {
+		b.WriteByte(']')
+		return
+	}
+
+	b.Grow(len(st) * (stackMinLen / 4))
+	st[0].format(b, s, verb)
+	for _, fr := range st[1:] {
+		b.WriteByte(' ')
+		fr.format(b, s, verb)
 	}
+	b.WriteByte(']')
 }
 
+// stackMinLen is a best-guess at the minimum length of a stack trace. It
+// doesn't need to be exact, just give a good enough head start for the buffer
+// to avoid the expensive early growth.
+const stackMinLen = 96
+
 // stack represents a stack of program counters.
 type stack []uintptr
 
@@ -136,10 +172,14 @@ func (s *stack) Format(st fmt.State, verb rune) {
 	case 'v':
 		switch {
 		case st.Flag('+'):
+			var b bytes.Buffer
+			b.Grow(len(*s) * stackMinLen)
 			for _, pc := range *s {
 				f := Frame(pc)
-				fmt.Fprintf(st, "\n%+v", f)
+				b.WriteByte('\n')
+				f.format(&b, st, 'v')
 			}
+			io.Copy(st, &b)
 		}
 	}
 }