perf: cache header order lookups instead of doing them per comparison

header.go

@@ -177,19 +177,31 @@ type HeaderKeyValues struct {
 type headerSorter struct {
 	kvs   []HeaderKeyValues
 	order map[string]int
+	// orderIdx[i], orderOK[i] cache order[strings.ToLower(kvs[i].Key)],
+	// resolved once per sort by SortedKeyValuesBy so that Less does no
+	// map lookups or lowercasing per comparison. Populated only when
+	// order is non-nil.
+	orderIdx []int
+	orderOK  []bool
 }
 
-func (s *headerSorter) Len() int      { return len(s.kvs) }
-func (s *headerSorter) Swap(i, j int) { s.kvs[i], s.kvs[j] = s.kvs[j], s.kvs[i] }
+func (s *headerSorter) Len() int { return len(s.kvs) }
+func (s *headerSorter) Swap(i, j int) {
+	s.kvs[i], s.kvs[j] = s.kvs[j], s.kvs[i]
+	// orderIdx/orderOK are only populated by SortedKeyValuesBy;
+	// SortedKeyValues sorts without them.
+	if s.order != nil {
+		s.orderIdx[i], s.orderIdx[j] = s.orderIdx[j], s.orderIdx[i]
+		s.orderOK[i], s.orderOK[j] = s.orderOK[j], s.orderOK[i]
+	}
+}
 func (s *headerSorter) Less(i, j int) bool {
 	// If the order isn't defined, sort lexicographically.
 	if s.order == nil {
 		return s.kvs[i].Key < s.kvs[j].Key
 	}
-	//idxi, iok := s.order[s.kvs[i].Key]
-	//idxj, jok := s.order[s.kvs[j].Key]
-	idxi, iok := s.order[strings.ToLower(s.kvs[i].Key)]
-	idxj, jok := s.order[strings.ToLower(s.kvs[j].Key)]
+	idxi, iok := s.orderIdx[i], s.orderOK[i]
+	idxj, jok := s.orderIdx[j], s.orderOK[j]
 	if !iok && !jok {
 		return s.kvs[i].Key < s.kvs[j].Key
 	} else if !iok && jok {
@@ -223,6 +235,10 @@ func (h Header) SortedKeyValues(exclude map[string]bool) (kvs []HeaderKeyValues,
 		mutex.RUnlock()
 	}
 	hs.kvs = kvs
+	// Reset any order left on the sorter by a previous SortedKeyValuesBy
+	// call, otherwise a pooled sorter sorts by the stale order instead of
+	// lexicographically.
+	hs.order = nil
 	sort.Sort(hs)
 	return kvs, hs
 }
@@ -242,6 +258,20 @@ func (h Header) SortedKeyValuesBy(order map[string]int, exclude map[string]bool)
 	}
 	hs.kvs = kvs
 	hs.order = order
+
+	// Decorate-sort-undecorate: resolve each key's order lookup once, so
+	// Less compares the cached results instead of doing two map lookups
+	// (with key lowercasing) per comparison.
+	if cap(hs.orderIdx) < len(kvs) {
+		hs.orderIdx = make([]int, len(kvs))
+		hs.orderOK = make([]bool, len(kvs))
+	}
+	hs.orderIdx = hs.orderIdx[:len(kvs)]
+	hs.orderOK = hs.orderOK[:len(kvs)]
+	for i, kv := range kvs {
+		hs.orderIdx[i], hs.orderOK[i] = order[strings.ToLower(kv.Key)]
+	}
+
 	sort.Sort(hs)
 
 	return kvs, hs

header_test.go

@@ -6,8 +6,10 @@ package http
 
 import (
 	"bytes"
+	"math/rand"
 	"reflect"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -359,3 +361,217 @@ func TestHTTP1HeaderOrder(t *testing.T) {
 		t.Fatalf("got:\n%swant:\n%s", buf.String(), expected)
 	}
 }
+
+func TestHeaderSorterPoolReuse(t *testing.T) {
+	// A sorter used by SortedKeyValuesBy keeps its order map. When it is
+	// pulled from the pool again by SortedKeyValues (no order), the stale
+	// order must not be used.
+	ordered := Header{
+		"Zebra": {"1"},
+		"Apple": {"2"},
+	}
+	// Deliberately the reverse of lexicographic order.
+	_, hs := ordered.SortedKeyValuesBy(map[string]int{"zebra": 0, "apple": 1}, nil)
+	headerSorterPool.Put(hs)
+
+	plain := Header{
+		"Apple":  {"1"},
+		"Banana": {"2"},
+		"Zebra":  {"3"},
+	}
+	kvs, _ := plain.SortedKeyValues(nil)
+	for i := 1; i < len(kvs); i++ {
+		if kvs[i-1].Key > kvs[i].Key {
+			t.Fatalf("keys not sorted lexicographically: %q before %q", kvs[i-1].Key, kvs[i].Key)
+		}
+	}
+}
+
+func TestSortedKeyValuesBy(t *testing.T) {
+	tests := []struct {
+		name  string
+		h     Header
+		order map[string]int
+		want  []string
+	}{
+		{
+			name: "all keys in order",
+			h: Header{
+				"Accept":     {"*/*"},
+				"User-Agent": {"x"},
+				"Referer":    {"y"},
+			},
+			order: map[string]int{"user-agent": 0, "referer": 1, "accept": 2},
+			want:  []string{"User-Agent", "Referer", "Accept"},
+		},
+		{
+			name: "keys absent from order sort lexicographically after ordered keys",
+			h: Header{
+				"Zeta":       {"z"},
+				"Alpha":      {"a"},
+				"Mid":        {"m"},
+				"In-Order":   {"x"},
+				"Also-Order": {"y"},
+			},
+			order: map[string]int{"in-order": 0, "also-order": 1},
+			want:  []string{"In-Order", "Also-Order", "Alpha", "Mid", "Zeta"},
+		},
+		{
+			name: "order lookup lowercases header keys",
+			h: Header{
+				"CONTENT-TYPE": {"a"},
+				"Accept":       {"b"},
+			},
+			order: map[string]int{"content-type": 0, "accept": 1},
+			want:  []string{"CONTENT-TYPE", "Accept"},
+		},
+		{
+			name: "no keys in order is fully lexicographic",
+			h: Header{
+				"B": {"1"},
+				"A": {"2"},
+				"C": {"3"},
+			},
+			order: map[string]int{"unrelated": 0},
+			want:  []string{"A", "B", "C"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			kvs, hs := tt.h.SortedKeyValuesBy(tt.order, nil)
+			got := make([]string, 0, len(kvs))
+			for _, kv := range kvs {
+				got = append(got, kv.Key)
+			}
+			headerSorterPool.Put(hs)
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Fatalf("SortedKeyValuesBy(%v) key order = %v, want %v", tt.order, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSortedKeyValuesByPoolReuse(t *testing.T) {
+	// Reuse one pooled sorter across ordered sorts of different sizes,
+	// with an orderless sort in between; per-sort state must be resized
+	// and repopulated on every call.
+	sortKeys := func(h Header, order map[string]int) []string {
+		var kvs []HeaderKeyValues
+		var hs *headerSorter
+		if order != nil {
+			kvs, hs = h.SortedKeyValuesBy(order, nil)
+		} else {
+			kvs, hs = h.SortedKeyValues(nil)
+		}
+		got := make([]string, 0, len(kvs))
+		for _, kv := range kvs {
+			got = append(got, kv.Key)
+		}
+		headerSorterPool.Put(hs)
+		return got
+	}
+
+	big := Header{"A": {"1"}, "B": {"2"}, "C": {"3"}, "D": {"4"}, "E": {"5"}}
+	bigOrder := map[string]int{"e": 0, "d": 1, "c": 2, "b": 3, "a": 4}
+	if got, want := sortKeys(big, bigOrder), []string{"E", "D", "C", "B", "A"}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("big ordered sort = %v, want %v", got, want)
+	}
+
+	small := Header{"Y": {"1"}, "X": {"2"}}
+	if got, want := sortKeys(small, map[string]int{"y": 0, "x": 1}), []string{"Y", "X"}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("small ordered sort after big = %v, want %v", got, want)
+	}
+
+	if got, want := sortKeys(big, nil), []string{"A", "B", "C", "D", "E"}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("orderless sort after ordered = %v, want %v", got, want)
+	}
+
+	if got, want := sortKeys(big, bigOrder), []string{"E", "D", "C", "B", "A"}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("ordered sort after orderless = %v, want %v", got, want)
+	}
+}
+
+func TestSortedKeyValuesByDuplicateOrderValues(t *testing.T) {
+	// A Header-Order: list with a repeated entry produces an order map
+	// whose values can reach or exceed len(order), e.g. ["c","a","c"]
+	// gives {"c": 2, "a": 1}. Keys present in the order map must still
+	// sort ahead of absent keys.
+	h := Header{
+		"Charlie": {"1"},
+		"Alpha":   {"2"},
+		"Mango":   {"3"}, // absent from order
+	}
+	order := map[string]int{"charlie": 2, "alpha": 1}
+	kvs, hs := h.SortedKeyValuesBy(order, nil)
+	got := make([]string, 0, len(kvs))
+	for _, kv := range kvs {
+		got = append(got, kv.Key)
+	}
+	headerSorterPool.Put(hs)
+	want := []string{"Alpha", "Charlie", "Mango"}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("SortedKeyValuesBy(%v) key order = %v, want %v", order, got, want)
+	}
+}
+
+// TestHeaderSorterLessEquivalence checks the cached-lookup Less against a
+// reference implementation of the original per-comparison semantics across
+// adversarial order maps: duplicate values, values at or past len(order),
+// negative values, and keys that collide when lowercased.
+func TestHeaderSorterLessEquivalence(t *testing.T) {
+	referenceLess := func(kvs []HeaderKeyValues, order map[string]int, i, j int) bool {
+		idxi, iok := order[strings.ToLower(kvs[i].Key)]
+		idxj, jok := order[strings.ToLower(kvs[j].Key)]
+		if !iok && !jok {
+			return kvs[i].Key < kvs[j].Key
+		} else if !iok && jok {
+			return false
+		} else if iok && !jok {
+			return true
+		}
+		return idxi < idxj
+	}
+
+	rng := rand.New(rand.NewSource(1))
+	keyPool := []string{
+		"Accept", "accept", "ACCEPT", "User-Agent", "user-agent",
+		"Cookie", "Referer", "X-A", "x-a", "Zeta", "alpha", "Alpha",
+	}
+	for iter := 0; iter < 200; iter++ {
+		n := 2 + rng.Intn(len(keyPool)-2)
+		keys := make([]string, n)
+		perm := rng.Perm(len(keyPool))
+		for i := range keys {
+			keys[i] = keyPool[perm[i]]
+		}
+
+		order := make(map[string]int)
+		for _, k := range keys {
+			if rng.Intn(2) == 0 {
+				order[strings.ToLower(k)] = rng.Intn(n+3) - 2 // gaps, duplicates, negatives
+			}
+		}
+
+		kvs := make([]HeaderKeyValues, n)
+		for i, k := range keys {
+			kvs[i] = HeaderKeyValues{Key: k, Values: []string{"v"}}
+		}
+
+		hs := &headerSorter{kvs: kvs, order: order}
+		hs.orderIdx = make([]int, n)
+		hs.orderOK = make([]bool, n)
+		for i, kv := range kvs {
+			hs.orderIdx[i], hs.orderOK[i] = order[strings.ToLower(kv.Key)]
+		}
+
+		for i := 0; i < n; i++ {
+			for j := 0; j < n; j++ {
+				if got, want := hs.Less(i, j), referenceLess(kvs, order, i, j); got != want {
+					t.Fatalf("iter %d: Less(%q, %q) with order %v = %v, want %v",
+						iter, kvs[i].Key, kvs[j].Key, order, got, want)
+				}
+			}
+		}
+	}
+}