Reuse header and control payload buffers

Author: nhooyr

header.go

@@ -75,12 +75,19 @@ func marshalHeader(h header) []byte {
 	return b
 }
 
+func makeHeaderBuf() []byte {
+	return make([]byte, maxHeaderSize-2)
+}
+
 // readHeader reads a header from the reader.
 // See https://tools.ietf.org/html/rfc6455#section-5.2
-func readHeader(r io.Reader) (header, error) {
-	// We read the first two bytes directly so that we know
+func readHeader(b []byte, r io.Reader) (header, error) {
+	if b == nil {
+		b = makeHeaderBuf()
+	}
+	// We read the first two bytes first so that we know
 	// exactly how long the header is.
-	b := make([]byte, 2, maxHeaderSize-2)
+	b = b[:2]
 	_, err := io.ReadFull(r, b)
 	if err != nil {
 		return header{}, err

header_test.go

@@ -32,7 +32,7 @@ func TestHeader(t *testing.T) {
 		b[2] |= 1 << 7
 
 		r := bytes.NewReader(b)
-		_, err := readHeader(r)
+		_, err := readHeader(nil, r)
 		if err == nil {
 			t.Fatalf("unexpected error value: %+v", err)
 		}
@@ -92,7 +92,7 @@ func TestHeader(t *testing.T) {
 func testHeader(t *testing.T, h header) {
 	b := marshalHeader(h)
 	r := bytes.NewReader(b)
-	h2, err := readHeader(r)
+	h2, err := readHeader(nil, r)
 	if err != nil {
 		t.Logf("header: %#v", h)
 		t.Logf("bytes: %b", b)

websocket.go

@@ -57,6 +57,9 @@ type Conn struct {
 
 	activePingsMu sync.Mutex
 	activePings   map[string]chan<- struct{}
+
+	headerBuf         []byte
+	controlPayloadBuf []byte
 }
 
 func (c *Conn) init() {
@@ -74,6 +77,9 @@ func (c *Conn) init() {
 
 	c.activePings = make(map[string]chan<- struct{})
 
+	c.headerBuf = makeHeaderBuf()
+	c.controlPayloadBuf = make([]byte, maxControlFramePayload)
+
 	runtime.SetFinalizer(c, func(c *Conn) {
 		c.close(xerrors.New("connection garbage collected"))
 	})
@@ -209,7 +215,7 @@ func (c *Conn) readFrameHeader(ctx context.Context) (header, error) {
 	case c.setReadTimeout <- ctx:
 	}
 
-	h, err := readHeader(c.br)
+	h, err := readHeader(c.headerBuf, c.br)
 	if err != nil {
 		select {
 		case <-c.closed:
@@ -249,8 +255,7 @@ func (c *Conn) handleControl(ctx context.Context, h header) error {
 	ctx, cancel := context.WithTimeout(ctx, time.Second*5)
 	defer cancel()
 
-	b := make([]byte, h.payloadLength)
-
+	b := c.controlPayloadBuf[:h.payloadLength]
 	_, err := c.readFramePayload(ctx, b)
 	if err != nil {
 		return err