Allow buffer reuse in ReadHeader/WriteHeader variants.

Author: iamtakingiteasy

read.go

@@ -21,6 +21,18 @@ func ReadHeader(r io.Reader) (h Header, err error) {
 	// So 14 - 2 = 12.
 	bts := make([]byte, 2, MaxHeaderSize-2)
 
+	return ReadHeaderBuffer(r, bts)
+}
+
+// ReadHeaderBuffer reads a frame header from r using user-provided buffer.
+// Provided buffer must be at least 12 bytes long.
+func ReadHeaderBuffer(r io.Reader, bts []byte) (h Header, err error) {
+	if cap(bts) < MaxHeaderSize-2 {
+		return h, io.ErrShortBuffer
+	}
+
+	bts = bts[:2]
+
 	// Prepare to hold first 2 bytes to choose size of next read.
 	_, err = io.ReadFull(r, bts)
 	if err != nil {

read_test.go

@@ -38,18 +38,40 @@ func TestReadHeader(t *testing.T) {
 }
 
 func BenchmarkReadHeader(b *testing.B) {
+	setup := func(header Header, n int) (rds []io.Reader) {
+		bts := MustCompileFrame(Frame{Header: header})
+		rds = make([]io.Reader, n)
+		for i := 0; i < n; i++ {
+			rds[i] = bytes.NewReader(bts)
+		}
+
+		return
+	}
+
 	for i, bench := range RWBenchCases {
 		b.Run(fmt.Sprintf("%s#%d", bench.label, i), func(b *testing.B) {
-			bts := MustCompileFrame(Frame{Header: bench.header})
-			rds := make([]io.Reader, b.N)
+			rds := setup(bench.header, b.N)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+
 			for i := 0; i < b.N; i++ {
-				rds[i] = bytes.NewReader(bts)
+				_, err := ReadHeader(rds[i])
+				if err != nil {
+					b.Fatal(err)
+				}
 			}
+		})
+
+		b.Run(fmt.Sprintf("reused-buffer-%s#%d", bench.label, i), func(b *testing.B) {
+			rds := setup(bench.header, b.N)
+			bts := make([]byte, MaxHeaderSize)
 
+			b.ReportAllocs()
 			b.ResetTimer()
 
 			for i := 0; i < b.N; i++ {
-				_, err := ReadHeader(rds[i])
+				_, err := ReadHeaderBuffer(rds[i], bts)
 				if err != nil {
 					b.Fatal(err)
 				}

write.go

@@ -50,11 +50,23 @@ func WriteHeader(w io.Writer, h Header) error {
 	// Make slice of bytes with capacity 14 that could hold any header.
 	bts := make([]byte, MaxHeaderSize)
 
+	return WriteHeaderBuffer(w, h, bts)
+}
+
+// WriteHeaderBuffer writes header binary representation into w using user-provided buffer.
+// Provided buffer must be at least 14 bytes long.
+func WriteHeaderBuffer(w io.Writer, h Header, bts []byte) error {
+	if cap(bts) < MaxHeaderSize {
+		return io.ErrShortBuffer
+	}
+
+	bts = bts[:MaxHeaderSize]
+
+	bts[0] = h.Rsv<<4 | byte(h.OpCode)
+
 	if h.Fin {
 		bts[0] |= bit0
 	}
-	bts[0] |= h.Rsv << 4
-	bts[0] |= byte(h.OpCode)
 
 	var n int
 	switch {

write_test.go

@@ -31,11 +31,27 @@ func TestWriteHeader(t *testing.T) {
 func BenchmarkWriteHeader(b *testing.B) {
 	for _, bench := range RWBenchCases {
 		b.Run(bench.label, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+
 			for i := 0; i < b.N; i++ {
 				if err := WriteHeader(ioutil.Discard, bench.header); err != nil {
 					b.Fatal(err)
 				}
 			}
 		})
+
+		b.Run("reused-buffer-"+bench.label, func(b *testing.B) {
+			bts := make([]byte, MaxHeaderSize)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				if err := WriteHeaderBuffer(ioutil.Discard, bench.header, bts); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
 	}
 }