net/http: speed up Header.WriteSubset

A few performance improvements, but without the stack sorting
change to avoid allocating, which is instead waiting on better
escape analysis.

R=rsc
CC=golang-dev
https://golang.org/cl/6265047

src/pkg/net/http/header.go

@@ -55,22 +55,54 @@ func (h Header) Write(w io.Writer) error {
 
 var headerNewlineToSpace = strings.NewReplacer("\n", " ", "\r", " ")
 
+type writeStringer interface {
+	WriteString(string) (int, error)
+}
+
+// stringWriter implements WriteString on a Writer.
+type stringWriter struct {
+	w io.Writer
+}
+
+func (w stringWriter) WriteString(s string) (n int, err error) {
+	return w.w.Write([]byte(s))
+}
+
+type keyValues struct {
+	key    string
+	values []string
+}
+
+type byKey []keyValues
+
+func (s byKey) Len() int           { return len(s) }
+func (s byKey) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+func (s byKey) Less(i, j int) bool { return s[i].key < s[j].key }
+
+func (h Header) sortedKeyValues(exclude map[string]bool) []keyValues {
+	kvs := make([]keyValues, 0, len(h))
+	for k, vv := range h {
+		if !exclude[k] {
+			kvs = append(kvs, keyValues{k, vv})
+		}
+	}
+	sort.Sort(byKey(kvs))
+	return kvs
+}
+
 // WriteSubset writes a header in wire format.
 // If exclude is not nil, keys where exclude[key] == true are not written.
 func (h Header) WriteSubset(w io.Writer, exclude map[string]bool) error {
-	keys := make([]string, 0, len(h))
-	for k := range h {
-		if !exclude[k] {
-			keys = append(keys, k)
-		}
+	ws, ok := w.(writeStringer)
+	if !ok {
+		ws = stringWriter{w}
 	}
-	sort.Strings(keys)
-	for _, k := range keys {
-		for _, v := range h[k] {
+	for _, kv := range h.sortedKeyValues(exclude) {
+		for _, v := range kv.values {
 			v = headerNewlineToSpace.Replace(v)
-			v = strings.TrimSpace(v)
-			for _, s := range []string{k, ": ", v, "\r\n"} {
-				if _, err := io.WriteString(w, s); err != nil {
+			v = textproto.TrimString(v)
+			for _, s := range []string{kv.key, ": ", v, "\r\n"} {
+				if _, err := ws.WriteString(s); err != nil {
 					return err
 				}
 			}

src/pkg/net/http/header_test.go

@@ -6,6 +6,7 @@ package http
 
 import (
 	"bytes"
+	"runtime"
 	"testing"
 )
 
@@ -67,6 +68,24 @@ var headerWriteTests = []struct {
 		nil,
 		"Blank: \r\nDouble-Blank: \r\nDouble-Blank: \r\n",
 	},
+	// Tests header sorting when over the insertion sort threshold side:
+	{
+		Header{
+			"k1": {"1a", "1b"},
+			"k2": {"2a", "2b"},
+			"k3": {"3a", "3b"},
+			"k4": {"4a", "4b"},
+			"k5": {"5a", "5b"},
+			"k6": {"6a", "6b"},
+			"k7": {"7a", "7b"},
+			"k8": {"8a", "8b"},
+			"k9": {"9a", "9b"},
+		},
+		map[string]bool{"k5": true},
+		"k1: 1a\r\nk1: 1b\r\nk2: 2a\r\nk2: 2b\r\nk3: 3a\r\nk3: 3b\r\n" +
+			"k4: 4a\r\nk4: 4b\r\nk6: 6a\r\nk6: 6b\r\n" +
+			"k7: 7a\r\nk7: 7b\r\nk8: 8a\r\nk8: 8b\r\nk9: 9a\r\nk9: 9b\r\n",
+	},
 }
 
 func TestHeaderWrite(t *testing.T) {
@@ -124,6 +143,18 @@ func TestHasToken(t *testing.T) {
 }
 
 func BenchmarkHeaderWriteSubset(b *testing.B) {
+	doHeaderWriteSubset(b.N, b)
+}
+
+func TestHeaderWriteSubsetMallocs(t *testing.T) {
+	doHeaderWriteSubset(100, t)
+}
+
+type errorfer interface {
+	Errorf(string, ...interface{})
+}
+
+func doHeaderWriteSubset(n int, t errorfer) {
 	h := Header(map[string][]string{
 		"Content-Length": {"123"},
 		"Content-Type":   {"text/plain"},
@@ -131,8 +162,17 @@ func BenchmarkHeaderWriteSubset(b *testing.B) {
 		"Server":         {"Go http package"},
 	})
 	var buf bytes.Buffer
-	for i := 0; i < b.N; i++ {
+	var m0 runtime.MemStats
+	runtime.ReadMemStats(&m0)
+	for i := 0; i < n; i++ {
 		buf.Reset()
 		h.WriteSubset(&buf, nil)
 	}
+	var m1 runtime.MemStats
+	runtime.ReadMemStats(&m1)
+	if mallocs := m1.Mallocs - m0.Mallocs; n >= 100 && mallocs >= uint64(n) {
+		// TODO(bradfitz,rsc): once we can sort with allocating,
+		// make this an error.  See http://golang.org/issue/3761
+		// t.Errorf("did %d mallocs (>= %d iterations); should have avoided mallocs", mallocs, n)
+	}
 }

src/pkg/net/textproto/textproto.go

@@ -121,3 +121,29 @@ func (c *Conn) Cmd(format string, args ...interface{}) (id uint, err error) {
 	}
 	return id, nil
 }
+
+// TrimString returns s without leading and trailing ASCII space.
+func TrimString(s string) string {
+	for len(s) > 0 && isASCIISpace(s[0]) {
+		s = s[1:]
+	}
+	for len(s) > 0 && isASCIISpace(s[len(s)-1]) {
+		s = s[:len(s)-1]
+	}
+	return s
+}
+
+// TrimBytes returns b without leading and trailing ASCII space.
+func TrimBytes(b []byte) []byte {
+	for len(b) > 0 && isASCIISpace(b[0]) {
+		b = b[1:]
+	}
+	for len(b) > 0 && isASCIISpace(b[len(b)-1]) {
+		b = b[:len(b)-1]
+	}
+	return b
+}
+
+func isASCIISpace(b byte) bool {
+	return b == ' ' || b == '\t' || b == '\n' || b == '\r'
+}