[IMPROVED] Reduce the size of the bufio.Reader for message header par…

…sing (nats-io#1187)
patrick-othmer · Jan 11, 2023 · bbde3f3 · bbde3f3
1 parent eada5ec
commit bbde3f3
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/nats.go b/nats.go
@@ -3559,7 +3559,8 @@ const (
 
 // decodeHeadersMsg will decode and headers.
 func decodeHeadersMsg(data []byte) (Header, error) {
-	tp := textproto.NewReader(bufio.NewReader(bytes.NewReader(data)))
+	br := bufio.NewReaderSize(bytes.NewReader(data), 128)
+	tp := textproto.NewReader(br)
 	l, err := tp.ReadLine()
 	if err != nil || len(l) < hdrPreEnd || l[:hdrPreEnd] != hdrLine[:hdrPreEnd] {
 		return nil, ErrBadHeaderMsg

diff --git a/nats_test.go b/nats_test.go
@@ -2967,3 +2967,46 @@ func TestServerListWithTrailingComma(t *testing.T) {
 		t.Fatalf("There should be only 1 URL in the list, got %v", l)
 	}
 }
+
+func BenchmarkHeaderDecode(b *testing.B) {
+	benchmarks := []struct {
+		name   string
+		header Header
+	}{
+		{"Small - 25", Header{
+			"Msg-ID": []string{"123"}},
+		},
+		{"Medium - 141", Header{
+			"CorrelationID": []string{"123"},
+			"Msg-ID":        []string{"456"},
+			"X-NATS-Keys":   []string{"A", "B", "C"},
+			"X-Test-Keys":   []string{"D", "E", "F"},
+		}},
+		{"Large - 368", Header{
+			"CorrelationID":     []string{"123"},
+			"Msg-ID":            []string{"456"},
+			"X-NATS-Keys":       []string{"A", "B", "C"},
+			"X-Test-Keys":       []string{"D", "E", "F"},
+			"X-A-Long-Header-1": []string{strings.Repeat("A", 100)},
+			"X-A-Long-Header-2": []string{strings.Repeat("A", 100)},
+		}},
+	}
+	for _, bm := range benchmarks {
+		b.Run(bm.name, func(b *testing.B) {
+			b.ReportAllocs()
+
+			m := NewMsg("foo")
+			m.Header = bm.header
+			hdr, err := m.headerBytes()
+			if err != nil {
+				b.Fatalf("Unexpected error: %v", err)
+			}
+
+			for i := 0; i < b.N; i++ {
+				if _, err := decodeHeadersMsg(hdr); err != nil {
+					b.Fatalf("Unexpected error: %v", err)
+				}
+			}
+		})
+	}
+}