Merge pull request #1 from jbub/encode

Optimize Attrs encoding by reducing allocations.

Author: jbub

.drone.yml

@@ -9,4 +9,4 @@ steps:
   - name: build
     image: golang:1.17
     commands:
-      - go test -race -v ./...
+      - go test -race -cover -v ./...

attrs.go

@@ -2,8 +2,6 @@ package sqlcommenter
 
 import (
 	"bytes"
-	"net/url"
-	"strings"
 )
 
 func AttrPairs(pairs ...string) Attrs {
@@ -39,25 +37,20 @@ func (a Attrs) encode(b *bytes.Buffer) {
 	sortKeys(keys)
 
 	for i, key := range keys {
-		b.WriteString(encodeKey(key))
+		writeQueryEscape(key, b)
+
 		b.WriteByte('=')
 		b.WriteByte('\'')
-		b.WriteString(encodeValue(a[key]))
+
+		writePathEscape(a[key], b)
+
 		b.WriteByte('\'')
 		if i < total-1 {
 			b.WriteByte(',')
 		}
 	}
 }
 
-func encodeKey(k string) string {
-	return url.QueryEscape(k)
-}
-
-func encodeValue(v string) string {
-	return strings.ReplaceAll(url.PathEscape(v), "+", "%20")
-}
-
 // sortKeys implements a simple insertion sort on string slice.
 // We save one alloc by not using sort.Strings.
 func sortKeys(keys []string) {

attrs_test.go

@@ -24,11 +24,11 @@ func TestAttrsEncode(t *testing.T) {
 		{
 			name: "multiple attrs",
 			attrs: map[string]string{
-				"key":  "value",
-				"2key": "value 33",
-				"key3": "44  value",
+				"key":  "DROP TABLE FOO",
+				"2key": "/param first",
+				"name": "1234",
 			},
-			want: "2key='value%2033',key='value',key3='44%20%20value'",
+			want: "2key='%2Fparam%20first',key='DROP%20TABLE%20FOO',name='1234'",
 		},
 	}
 

comment.go

@@ -75,6 +75,6 @@ func (c *commenter) attrs(ctx context.Context) Attrs {
 
 var bufPool = sync.Pool{
 	New: func() interface{} {
-		return &bytes.Buffer{}
+		return bytes.NewBuffer(make([]byte, 0, 100))
 	},
 }

comment_test.go

@@ -77,16 +77,17 @@ func TestCommentConcurrent(t *testing.T) {
 }
 
 func BenchmarkComment(b *testing.B) {
-	b.ReportAllocs()
-	b.SetBytes(2)
-
 	ctx := context.Background()
 	cmt := newCommenter(WithAttrs(map[string]string{
 		"key":  "value",
 		"2key": "value 33",
 		"key3": "44  value",
 	}))
 
+	b.ReportAllocs()
+	b.SetBytes(2)
+	b.ResetTimer()
+
 	for i := 0; i < b.N; i++ {
 		cmt.comment(ctx, "SELECT * FROM my_table WHERE column IS NOT NULL")
 	}

escape.go

@@ -0,0 +1,77 @@
+package sqlcommenter
+
+// Code is adapted from standard library package net/url.
+// Copyright (c) 2009 The Go Authors. All rights reserved.
+
+import (
+	"bytes"
+)
+
+const upperhex = "0123456789ABCDEF"
+
+func writeQueryEscape(s string, b *bytes.Buffer) {
+	writeEscape(s, true, b)
+}
+
+func writePathEscape(s string, b *bytes.Buffer) {
+	writeEscape(s, false, b)
+}
+
+func writeEscape(s string, query bool, b *bytes.Buffer) {
+	spaceCount, hexCount := 0, 0
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		if shouldEscape(c, query) {
+			if c == ' ' && query {
+				spaceCount++
+			} else {
+				hexCount++
+			}
+		}
+	}
+
+	if spaceCount == 0 && hexCount == 0 {
+		b.WriteString(s)
+		return
+	}
+
+	if hexCount == 0 {
+		for i := 0; i < len(s); i++ {
+			if s[i] == ' ' {
+				b.WriteByte('+')
+			} else {
+				b.WriteByte(s[i])
+			}
+		}
+		return
+	}
+
+	for i := 0; i < len(s); i++ {
+		switch c := s[i]; {
+		case c == ' ' && query:
+			b.WriteByte('+')
+		case shouldEscape(c, query):
+			b.WriteByte('%')
+			b.WriteByte(upperhex[c>>4])
+			b.WriteByte(upperhex[c&15])
+		default:
+			b.WriteByte(c)
+		}
+	}
+}
+
+func shouldEscape(c byte, query bool) bool {
+	if 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || '0' <= c && c <= '9' {
+		return false
+	}
+	switch c {
+	case '-', '_', '.', '~':
+		return false
+	case '$', '&', '+', ',', '/', ':', ';', '=', '?', '@':
+		if query {
+			return true
+		}
+		return c == '/' || c == ';' || c == ',' || c == '?'
+	}
+	return true
+}

escape_test.go

@@ -0,0 +1,93 @@
+package sqlcommenter
+
+// Code is adapted from standard library package net/url.
+// Copyright (c) 2009 The Go Authors. All rights reserved.
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestQueryEscape(t *testing.T) {
+	cases := []struct {
+		in   string
+		want string
+	}{
+		{
+			in:   "",
+			want: "",
+		},
+		{
+			in:   "abc",
+			want: "abc",
+		},
+		{
+			in:   "one two",
+			want: "one+two",
+		},
+		{
+			in:   "10%",
+			want: "10%25",
+		},
+		{
+			in:   " ?&=#+%!<>#\"{}|\\^[]`☺\t:/@$'()*,;",
+			want: "+%3F%26%3D%23%2B%25%21%3C%3E%23%22%7B%7D%7C%5C%5E%5B%5D%60%E2%98%BA%09%3A%2F%40%24%27%28%29%2A%2C%3B",
+		},
+	}
+
+	for _, cs := range cases {
+		t.Run(cs.in, func(t *testing.T) {
+			var b bytes.Buffer
+			writeQueryEscape(cs.in, &b)
+			if got := b.String(); cs.want != got {
+				t.Errorf("got %q, want %q", got, cs.want)
+			}
+		})
+	}
+}
+
+func TestPathEscape(t *testing.T) {
+	cases := []struct {
+		in   string
+		want string
+	}{
+		{
+			in:   "",
+			want: "",
+		},
+		{
+			in:   "abc",
+			want: "abc",
+		},
+		{
+			in:   "abc+def",
+			want: "abc+def",
+		},
+		{
+			in:   "a/b",
+			want: "a%2Fb",
+		},
+		{
+			in:   "one two",
+			want: "one%20two",
+		},
+		{
+			in:   "10%",
+			want: "10%25",
+		},
+		{
+			in:   " ?&=#+%!<>#\"{}|\\^[]`☺\t:/@$'()*,;",
+			want: "%20%3F&=%23+%25%21%3C%3E%23%22%7B%7D%7C%5C%5E%5B%5D%60%E2%98%BA%09:%2F@$%27%28%29%2A%2C%3B",
+		},
+	}
+
+	for _, cs := range cases {
+		t.Run(cs.in, func(t *testing.T) {
+			var b bytes.Buffer
+			writePathEscape(cs.in, &b)
+			if got := b.String(); cs.want != got {
+				t.Errorf("got %q, want %q", got, cs.want)
+			}
+		})
+	}
+}