From a2af651ecb5565c495d8f670ee740045d188dd1e Mon Sep 17 00:00:00 2001
From: Charlie Vieth <charlie.vieth@gmail.com>
Date: Wed, 17 Jul 2019 02:57:46 -0400
Subject: [PATCH] replacer: evaluate format once and improve perf by ~3x
 (#3002)

* replacer: evaluate format once and improve perf by ~3x

This improves the performance of logging by almost 3x and reduces memory
usage by ~8x.

Benchmark results:

benchmark                                old ns/op     new ns/op     delta
BenchmarkReplacer-12                     644           324           -49.69%
BenchmarkReplacer_CommonLogFormat-12     4228          1471          -65.21%

benchmark                                old allocs     new allocs     delta
BenchmarkReplacer-12                     8              2              -75.00%
BenchmarkReplacer_CommonLogFormat-12     51             17             -66.67%

benchmark                                old bytes     new bytes     delta
BenchmarkReplacer-12                     240           48            -80.00%
BenchmarkReplacer_CommonLogFormat-12     3723          446           -88.02%

* replacer: code review comments

* bufPool: document why we use a pointer to a slice
* parseFormat: fix confusing comment
* TestParseFormat_Nodes: rename to TestParseFormatNodes

* replacer: use a value for bufPool instead of a pointer

* replacer: remove comment

* replacer: replace labels with knownLabels

The previous slice of label names is no longer needed.
---
 plugin/pkg/replacer/replacer.go      | 426 ++++++++++++++++-----------
 plugin/pkg/replacer/replacer_test.go | 207 ++++++++++++-
 2 files changed, 453 insertions(+), 180 deletions(-)

diff --git a/plugin/pkg/replacer/replacer.go b/plugin/pkg/replacer/replacer.go
index 3e2cbfcbd..81261ca79 100644
--- a/plugin/pkg/replacer/replacer.go
+++ b/plugin/pkg/replacer/replacer.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/coredns/coredns/plugin/metadata"
@@ -14,195 +15,262 @@ import (
 )
 
 // Replacer replaces labels for values in strings.
-type Replacer struct {
-	valueFunc func(request.Request, *dnstest.Recorder, string) string
-	labels    []string
-}
+type Replacer struct{}
 
-// labels are all supported labels that can be used in the default Replacer.
-var labels = []string{
-	"{type}",
-	"{name}",
-	"{class}",
-	"{proto}",
-	"{size}",
-	"{remote}",
-	"{port}",
-	"{local}",
-	// Header values.
-	headerReplacer + "id}",
-	headerReplacer + "opcode}",
-	headerReplacer + "do}",
-	headerReplacer + "bufsize}",
-	// Recorded replacements.
-	"{rcode}",
-	"{rsize}",
-	"{duration}",
-	headerReplacer + "rflags}",
-}
-
-// value returns the current value of label.
-func value(state request.Request, rr *dnstest.Recorder, label string) string {
-	switch label {
-	case "{type}":
-		return state.Type()
-	case "{name}":
-		return state.Name()
-	case "{class}":
-		return state.Class()
-	case "{proto}":
-		return state.Proto()
-	case "{size}":
-		return strconv.Itoa(state.Req.Len())
-	case "{remote}":
-		return addrToRFC3986(state.IP())
-	case "{port}":
-		return state.Port()
-	case "{local}":
-		return addrToRFC3986(state.LocalIP())
-	// Header placeholders (case-insensitive).
-	case headerReplacer + "id}":
-		return strconv.Itoa(int(state.Req.Id))
-	case headerReplacer + "opcode}":
-		return strconv.Itoa(state.Req.Opcode)
-	case headerReplacer + "do}":
-		return boolToString(state.Do())
-	case headerReplacer + "bufsize}":
-		return strconv.Itoa(state.Size())
-	// Recorded replacements.
-	case "{rcode}":
-		if rr == nil {
-			return EmptyValue
-		}
-		rcode := dns.RcodeToString[rr.Rcode]
-		if rcode == "" {
-			rcode = strconv.Itoa(rr.Rcode)
-		}
-		return rcode
-	case "{rsize}":
-		if rr == nil {
-			return EmptyValue
-		}
-		return strconv.Itoa(rr.Len)
-	case "{duration}":
-		if rr == nil {
-			return EmptyValue
-		}
-		return strconv.FormatFloat(time.Since(rr.Start).Seconds(), 'f', -1, 64) + "s"
-	case headerReplacer + "rflags}":
-		if rr != nil && rr.Msg != nil {
-			return flagsToString(rr.Msg.MsgHdr)
-		}
-		return EmptyValue
-	}
-	return EmptyValue
-}
-
-// New makes a new replacer. This only needs to be called once in the setup and then call Replace for each incoming message.
-// A replacer is safe for concurrent use.
+// New makes a new replacer. This only needs to be called once in the setup and
+// then call Replace for each incoming message. A replacer is safe for concurrent use.
 func New() Replacer {
-	return Replacer{
-		valueFunc: value,
-		labels:    labels,
-	}
+	return Replacer{}
 }
 
 // Replace performs a replacement of values on s and returns the string with the replaced values.
 func (r Replacer) Replace(ctx context.Context, state request.Request, rr *dnstest.Recorder, s string) string {
-	for _, placeholder := range r.labels {
-		if strings.Contains(s, placeholder) {
-			s = strings.Replace(s, placeholder, r.valueFunc(state, rr, placeholder), -1)
-		}
-	}
-
-	// Metadata label replacements. Scan for {/ and search for next }, replace that metadata label with
-	// any meta data that is available.
-	b := strings.Builder{}
-	for strings.Contains(s, labelReplacer) {
-		idxStart := strings.Index(s, labelReplacer)
-		endOffset := idxStart + len(labelReplacer)
-		idxEnd := strings.Index(s[endOffset:], "}")
-		if idxEnd > -1 {
-			label := s[idxStart+2 : endOffset+idxEnd]
-
-			fm := metadata.ValueFunc(ctx, label)
-			replacement := EmptyValue
-			if fm != nil {
-				replacement = fm()
-			}
-
-			b.WriteString(s[:idxStart])
-			b.WriteString(replacement)
-			s = s[endOffset+idxEnd+1:]
-		} else {
-			break
-		}
-	}
-
-	b.WriteString(s)
-	return b.String()
-}
-
-func boolToString(b bool) string {
-	if b {
-		return "true"
-	}
-	return "false"
-}
-
-// flagsToString checks all header flags and returns those
-// that are set as a string separated with commas
-func flagsToString(h dns.MsgHdr) string {
-	flags := make([]string, 7)
-	i := 0
-
-	if h.Response {
-		flags[i] = "qr"
-		i++
-	}
-
-	if h.Authoritative {
-		flags[i] = "aa"
-		i++
-	}
-	if h.Truncated {
-		flags[i] = "tc"
-		i++
-	}
-	if h.RecursionDesired {
-		flags[i] = "rd"
-		i++
-	}
-	if h.RecursionAvailable {
-		flags[i] = "ra"
-		i++
-	}
-	if h.Zero {
-		flags[i] = "z"
-		i++
-	}
-	if h.AuthenticatedData {
-		flags[i] = "ad"
-		i++
-	}
-	if h.CheckingDisabled {
-		flags[i] = "cd"
-		i++
-	}
-	return strings.Join(flags[:i], ",")
-}
-
-// addrToRFC3986 will add brackets to the address if it is an IPv6 address.
-func addrToRFC3986(addr string) string {
-	if strings.Contains(addr, ":") {
-		return "[" + addr + "]"
-	}
-	return addr
+	return loadFormat(s).Replace(ctx, state, rr)
 }
 
 const (
 	headerReplacer = "{>"
-	labelReplacer  = "{/"
 	// EmptyValue is the default empty value.
 	EmptyValue = "-"
 )
+
+// labels are all supported labels that can be used in the default Replacer.
+var labels = map[string]struct{}{
+	"{type}":   {},
+	"{name}":   {},
+	"{class}":  {},
+	"{proto}":  {},
+	"{size}":   {},
+	"{remote}": {},
+	"{port}":   {},
+	"{local}":  {},
+	// Header values.
+	headerReplacer + "id}":      {},
+	headerReplacer + "opcode}":  {},
+	headerReplacer + "do}":      {},
+	headerReplacer + "bufsize}": {},
+	// Recorded replacements.
+	"{rcode}":                  {},
+	"{rsize}":                  {},
+	"{duration}":               {},
+	headerReplacer + "rflags}": {},
+}
+
+// appendValue appends the current value of label.
+func appendValue(b []byte, state request.Request, rr *dnstest.Recorder, label string) []byte {
+	switch label {
+	case "{type}":
+		return append(b, state.Type()...)
+	case "{name}":
+		return append(b, state.Name()...)
+	case "{class}":
+		return append(b, state.Class()...)
+	case "{proto}":
+		return append(b, state.Proto()...)
+	case "{size}":
+		return strconv.AppendInt(b, int64(state.Req.Len()), 10)
+	case "{remote}":
+		return appendAddrToRFC3986(b, state.IP())
+	case "{port}":
+		return append(b, state.Port()...)
+	case "{local}":
+		return appendAddrToRFC3986(b, state.LocalIP())
+	// Header placeholders (case-insensitive).
+	case headerReplacer + "id}":
+		return strconv.AppendInt(b, int64(state.Req.Id), 10)
+	case headerReplacer + "opcode}":
+		return strconv.AppendInt(b, int64(state.Req.Opcode), 10)
+	case headerReplacer + "do}":
+		return strconv.AppendBool(b, state.Do())
+	case headerReplacer + "bufsize}":
+		return strconv.AppendInt(b, int64(state.Size()), 10)
+	// Recorded replacements.
+	case "{rcode}":
+		if rr == nil {
+			return append(b, EmptyValue...)
+		}
+		if rcode := dns.RcodeToString[rr.Rcode]; rcode != "" {
+			return append(b, rcode...)
+		}
+		return strconv.AppendInt(b, int64(rr.Rcode), 10)
+	case "{rsize}":
+		if rr == nil {
+			return append(b, EmptyValue...)
+		}
+		return strconv.AppendInt(b, int64(rr.Len), 10)
+	case "{duration}":
+		if rr == nil {
+			return append(b, EmptyValue...)
+		}
+		secs := time.Since(rr.Start).Seconds()
+		return append(strconv.AppendFloat(b, secs, 'f', -1, 64), 's')
+	case headerReplacer + "rflags}":
+		if rr != nil && rr.Msg != nil {
+			return appendFlags(b, rr.Msg.MsgHdr)
+		}
+		return append(b, EmptyValue...)
+	default:
+		return append(b, EmptyValue...)
+	}
+}
+
+// appendFlags checks all header flags and appends those
+// that are set as a string separated with commas
+func appendFlags(b []byte, h dns.MsgHdr) []byte {
+	origLen := len(b)
+	if h.Response {
+		b = append(b, "qr,"...)
+	}
+	if h.Authoritative {
+		b = append(b, "aa,"...)
+	}
+	if h.Truncated {
+		b = append(b, "tc,"...)
+	}
+	if h.RecursionDesired {
+		b = append(b, "rd,"...)
+	}
+	if h.RecursionAvailable {
+		b = append(b, "ra,"...)
+	}
+	if h.Zero {
+		b = append(b, "z,"...)
+	}
+	if h.AuthenticatedData {
+		b = append(b, "ad,"...)
+	}
+	if h.CheckingDisabled {
+		b = append(b, "cd,"...)
+	}
+	if n := len(b); n > origLen {
+		return b[:n-1] // trim trailing ','
+	}
+	return b
+}
+
+// appendAddrToRFC3986 will add brackets to the address if it is an IPv6 address.
+func appendAddrToRFC3986(b []byte, addr string) []byte {
+	if strings.IndexByte(addr, ':') != -1 {
+		b = append(b, '[')
+		b = append(b, addr...)
+		b = append(b, ']')
+	} else {
+		b = append(b, addr...)
+	}
+	return b
+}
+
+type nodeType int
+
+const (
+	typeLabel    nodeType = iota // "{type}"
+	typeLiteral                  // "foo"
+	typeMetadata                 // "{/metadata}"
+)
+
+// A node represents a segment of a parsed format.  For example: "A {type}"
+// contains two nodes: "A " (literal); and "{type}" (label).
+type node struct {
+	value string // Literal value, label or metadata label
+	typ   nodeType
+}
+
+// A replacer is an ordered list of all the nodes in a format.
+type replacer []node
+
+func parseFormat(s string) replacer {
+	// Assume there is a literal between each label - its cheaper to over
+	// allocate once than allocate twice.
+	rep := make(replacer, 0, strings.Count(s, "{")*2)
+	for {
+		// We find the right bracket then backtrack to find the left bracket.
+		// This allows us to handle formats like: "{ {foo} }".
+		j := strings.IndexByte(s, '}')
+		if j < 0 {
+			break
+		}
+		i := strings.LastIndexByte(s[:j], '{')
+		if i < 0 {
+			// Handle: "A } {foo}" by treating "A }" as a literal
+			rep = append(rep, node{
+				value: s[:j+1],
+				typ:   typeLiteral,
+			})
+			s = s[j+1:]
+			continue
+		}
+
+		val := s[i : j+1]
+		var typ nodeType
+		switch _, ok := labels[val]; {
+		case ok:
+			typ = typeLabel
+		case strings.HasPrefix(val, "{/"):
+			// Strip "{/}" from metadata labels
+			val = val[2 : len(val)-1]
+			typ = typeMetadata
+		default:
+			// Given: "A {X}" val is "{X}" expand it to the whole literal.
+			val = s[:j+1]
+			typ = typeLiteral
+		}
+
+		// Append any leading literal.  Given "A {type}" the literal is "A "
+		if i != 0 && typ != typeLiteral {
+			rep = append(rep, node{
+				value: s[:i],
+				typ:   typeLiteral,
+			})
+		}
+		rep = append(rep, node{
+			value: val,
+			typ:   typ,
+		})
+		s = s[j+1:]
+	}
+	if len(s) != 0 {
+		rep = append(rep, node{
+			value: s,
+			typ:   typeLiteral,
+		})
+	}
+	return rep
+}
+
+var replacerCache sync.Map // map[string]replacer
+
+func loadFormat(s string) replacer {
+	if v, ok := replacerCache.Load(s); ok {
+		return v.(replacer)
+	}
+	v, _ := replacerCache.LoadOrStore(s, parseFormat(s))
+	return v.(replacer)
+}
+
+// bufPool stores pointers to scratch buffers.
+var bufPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, 0, 256)
+	},
+}
+
+func (r replacer) Replace(ctx context.Context, state request.Request, rr *dnstest.Recorder) string {
+	b := bufPool.Get().([]byte)
+	for _, s := range r {
+		switch s.typ {
+		case typeLabel:
+			b = appendValue(b, state, rr, s.value)
+		case typeLiteral:
+			b = append(b, s.value...)
+		case typeMetadata:
+			if fm := metadata.ValueFunc(ctx, s.value); fm != nil {
+				b = append(b, fm()...)
+			} else {
+				b = append(b, EmptyValue...)
+			}
+		}
+	}
+	s := string(b)
+	bufPool.Put(b[:0])
+	return s
+}
diff --git a/plugin/pkg/replacer/replacer_test.go b/plugin/pkg/replacer/replacer_test.go
index 13ba396b9..3e0a50691 100644
--- a/plugin/pkg/replacer/replacer_test.go
+++ b/plugin/pkg/replacer/replacer_test.go
@@ -2,6 +2,8 @@ package replacer
 
 import (
 	"context"
+	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/coredns/coredns/plugin/metadata"
@@ -12,6 +14,9 @@ import (
 	"github.com/miekg/dns"
 )
 
+// This is the default format used by the log package
+const CommonLogFormat = `{remote}:{port} - {>id} "{type} {class} {name} {proto} {size} {>do} {>bufsize}" {rcode} {>rflags} {rsize} {duration}`
+
 func TestReplacer(t *testing.T) {
 	w := dnstest.NewRecorder(&test.ResponseWriter{})
 	r := new(dns.Msg)
@@ -32,6 +37,178 @@ func TestReplacer(t *testing.T) {
 	}
 }
 
+func TestParseFormat(t *testing.T) {
+	type formatTest struct {
+		Format   string
+		Expected replacer
+	}
+	tests := []formatTest{
+		{
+			Format:   "",
+			Expected: replacer{},
+		},
+		{
+			Format: "A",
+			Expected: replacer{
+				{"A", typeLiteral},
+			},
+		},
+		{
+			Format: "A {A}",
+			Expected: replacer{
+				{"A {A}", typeLiteral},
+			},
+		},
+		{
+			Format: "{{remote}}",
+			Expected: replacer{
+				{"{", typeLiteral},
+				{"{remote}", typeLabel},
+				{"}", typeLiteral},
+			},
+		},
+		{
+			Format: "{ A {remote} A }",
+			Expected: replacer{
+				{"{ A ", typeLiteral},
+				{"{remote}", typeLabel},
+				{" A }", typeLiteral},
+			},
+		},
+		{
+			Format: "{remote}}",
+			Expected: replacer{
+				{"{remote}", typeLabel},
+				{"}", typeLiteral},
+			},
+		},
+		{
+			Format: "{{remote}",
+			Expected: replacer{
+				{"{", typeLiteral},
+				{"{remote}", typeLabel},
+			},
+		},
+		{
+			Format: `Foo } {remote}`,
+			Expected: replacer{
+				// we don't do any optimizations to join adjacent literals
+				{"Foo }", typeLiteral},
+				{" ", typeLiteral},
+				{"{remote}", typeLabel},
+			},
+		},
+		{
+			Format: `{ Foo`,
+			Expected: replacer{
+				{"{ Foo", typeLiteral},
+			},
+		},
+		{
+			Format: `} Foo`,
+			Expected: replacer{
+				{"}", typeLiteral},
+				{" Foo", typeLiteral},
+			},
+		},
+		{
+			Format: "A { {remote} {type} {/meta1} } B",
+			Expected: replacer{
+				{"A { ", typeLiteral},
+				{"{remote}", typeLabel},
+				{" ", typeLiteral},
+				{"{type}", typeLabel},
+				{" ", typeLiteral},
+				{"meta1", typeMetadata},
+				{" }", typeLiteral},
+				{" B", typeLiteral},
+			},
+		},
+		{
+			Format: `LOG {remote}:{port} - {>id} "{type} {class} {name} {proto} ` +
+				`{size} {>do} {>bufsize}" {rcode} {>rflags} {rsize} {/meta1}-{/meta2} ` +
+				`{duration} END OF LINE`,
+			Expected: replacer{
+				{"LOG ", typeLiteral},
+				{"{remote}", typeLabel},
+				{":", typeLiteral},
+				{"{port}", typeLabel},
+				{" - ", typeLiteral},
+				{"{>id}", typeLabel},
+				{` "`, typeLiteral},
+				{"{type}", typeLabel},
+				{" ", typeLiteral},
+				{"{class}", typeLabel},
+				{" ", typeLiteral},
+				{"{name}", typeLabel},
+				{" ", typeLiteral},
+				{"{proto}", typeLabel},
+				{" ", typeLiteral},
+				{"{size}", typeLabel},
+				{" ", typeLiteral},
+				{"{>do}", typeLabel},
+				{" ", typeLiteral},
+				{"{>bufsize}", typeLabel},
+				{`" `, typeLiteral},
+				{"{rcode}", typeLabel},
+				{" ", typeLiteral},
+				{"{>rflags}", typeLabel},
+				{" ", typeLiteral},
+				{"{rsize}", typeLabel},
+				{" ", typeLiteral},
+				{"meta1", typeMetadata},
+				{"-", typeLiteral},
+				{"meta2", typeMetadata},
+				{" ", typeLiteral},
+				{"{duration}", typeLabel},
+				{" END OF LINE", typeLiteral},
+			},
+		},
+	}
+	for i, x := range tests {
+		r := parseFormat(x.Format)
+		if !reflect.DeepEqual(r, x.Expected) {
+			t.Errorf("%d: Expected:\n\t%+v\nGot:\n\t%+v", i, x.Expected, r)
+		}
+	}
+}
+
+func TestParseFormatNodes(t *testing.T) {
+	// If we parse the format successfully the result of joining all the
+	// segments should match the original format.
+	formats := []string{
+		"",
+		"msg",
+		"{remote}",
+		"{remote}",
+		"{{remote}",
+		"{{remote}}",
+		"{{remote}} A",
+		CommonLogFormat,
+		CommonLogFormat + " FOO} {BAR}",
+		"A " + CommonLogFormat + " FOO} {BAR}",
+		"A " + CommonLogFormat + " {/meta}",
+	}
+	join := func(r replacer) string {
+		a := make([]string, len(r))
+		for i, n := range r {
+			if n.typ == typeMetadata {
+				a[i] = "{/" + n.value + "}"
+			} else {
+				a[i] = n.value
+			}
+		}
+		return strings.Join(a, "")
+	}
+	for _, format := range formats {
+		r := parseFormat(format)
+		s := join(r)
+		if s != format {
+			t.Errorf("Expected format to be: '%s' got: '%s'", format, s)
+		}
+	}
+}
+
 func TestLabels(t *testing.T) {
 	w := dnstest.NewRecorder(&test.ResponseWriter{})
 	r := new(dns.Msg)
@@ -68,7 +245,7 @@ func TestLabels(t *testing.T) {
 		t.Fatalf("Expect %d labels, got %d", len(expect), len(labels))
 	}
 
-	for _, lbl := range labels {
+	for lbl := range labels {
 		repl := replacer.Replace(ctx, state, w, lbl)
 		if lbl == "{duration}" {
 			if repl[len(repl)-1] != 's' {
@@ -98,6 +275,34 @@ func BenchmarkReplacer(b *testing.B) {
 	}
 }
 
+func BenchmarkReplacer_CommonLogFormat(b *testing.B) {
+
+	w := dnstest.NewRecorder(&test.ResponseWriter{})
+	r := new(dns.Msg)
+	r.SetQuestion("example.org.", dns.TypeHINFO)
+	r.Id = 1053
+	r.AuthenticatedData = true
+	r.CheckingDisabled = true
+	r.MsgHdr.AuthenticatedData = true
+	w.WriteMsg(r)
+	state := request.Request{W: w, Req: r}
+
+	replacer := New()
+	ctxt := context.TODO()
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		replacer.Replace(ctxt, state, w, CommonLogFormat)
+	}
+}
+
+func BenchmarkParseFormat(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		parseFormat(CommonLogFormat)
+	}
+}
+
 type testProvider map[string]metadata.Func
 
 func (tp testProvider) Metadata(ctx context.Context, state request.Request) context.Context {