Pool buffers in wspb and wsjson

Closes #71

Author: nhooyr

internal/bpool/bpool.go

@@ -0,0 +1,24 @@
+package bpool
+
+import (
+	"bytes"
+	"sync"
+)
+
+var bpool sync.Pool
+
+// Get returns a buffer from the pool or creates a new one if
+// the pool is empty.
+func Get() *bytes.Buffer {
+	b, ok := bpool.Get().(*bytes.Buffer)
+	if !ok {
+		return bytes.NewBuffer(make([]byte, 0, 4096))
+	}
+	return b
+}
+
+// Put returns a buffer into the pool.
+func Put(b *bytes.Buffer) {
+	b.Reset()
+	bpool.Put(b)
+}

internal/bpool/bpool_test.go

@@ -0,0 +1,47 @@
+package bpool
+
+import (
+	"strconv"
+	"sync"
+	"testing"
+)
+
+func BenchmarkSyncPool(b *testing.B) {
+	sizes := []int{
+		2,
+		16,
+		32,
+		64,
+		128,
+		256,
+		512,
+		4096,
+		16384,
+	}
+	for _, size := range sizes {
+		b.Run(strconv.Itoa(size), func(b *testing.B) {
+			b.Run("allocate", func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					buf := make([]byte, size)
+					_ = buf
+				}
+			})
+			b.Run("pool", func(b *testing.B) {
+				b.ReportAllocs()
+
+				p := sync.Pool{}
+
+				b.ResetTimer()
+				for i := 0; i < b.N; i++ {
+					buf := p.Get()
+					if buf == nil {
+						buf = make([]byte, size)
+					}
+
+					p.Put(buf)
+				}
+			})
+		})
+	}
+}

wsjson/wsjson.go

@@ -4,10 +4,12 @@ package wsjson
 import (
 	"context"
 	"encoding/json"
+	"sync"
 
 	"golang.org/x/xerrors"
 
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/internal/bpool"
 )
 
 // Read reads a json message from c into v.
@@ -21,8 +23,10 @@ func Read(ctx context.Context, c *websocket.Conn, v interface{}) error {
 	return nil
 }
 
+var bufPool sync.Pool
+
 func read(ctx context.Context, c *websocket.Conn, v interface{}) error {
-	typ, b, err := c.Read(ctx)
+	typ, r, err := c.Reader(ctx)
 	if err != nil {
 		return err
 	}
@@ -32,7 +36,17 @@ func read(ctx context.Context, c *websocket.Conn, v interface{}) error {
 		return xerrors.Errorf("unexpected frame type for json (expected %v): %v", websocket.MessageText, typ)
 	}
 
-	err = json.Unmarshal(b, v)
+	b := bpool.Get()
+	defer func() {
+		bpool.Put(b)
+	}()
+
+	_, err = b.ReadFrom(r)
+	if err != nil {
+		return err
+	}
+
+	err = json.Unmarshal(b.Bytes(), v)
 	if err != nil {
 		return xerrors.Errorf("failed to unmarshal json: %w", err)
 	}

wspb/wspb.go

@@ -3,11 +3,13 @@ package wspb
 
 import (
 	"context"
+	"sync"
 
 	"github.com/golang/protobuf/proto"
 	"golang.org/x/xerrors"
 
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/internal/bpool"
 )
 
 // Read reads a protobuf message from c into v.
@@ -21,7 +23,7 @@ func Read(ctx context.Context, c *websocket.Conn, v proto.Message) error {
 }
 
 func read(ctx context.Context, c *websocket.Conn, v proto.Message) error {
-	typ, b, err := c.Read(ctx)
+	typ, r, err := c.Reader(ctx)
 	if err != nil {
 		return err
 	}
@@ -31,7 +33,17 @@ func read(ctx context.Context, c *websocket.Conn, v proto.Message) error {
 		return xerrors.Errorf("unexpected frame type for protobuf (expected %v): %v", websocket.MessageBinary, typ)
 	}
 
-	err = proto.Unmarshal(b, v)
+	b := bpool.Get()
+	defer func() {
+		bpool.Put(b)
+	}()
+
+	_, err = b.ReadFrom(r)
+	if err != nil {
+		return err
+	}
+
+	err = proto.Unmarshal(b.Bytes(), v)
 	if err != nil {
 		return xerrors.Errorf("failed to unmarshal protobuf: %w", err)
 	}
@@ -49,11 +61,22 @@ func Write(ctx context.Context, c *websocket.Conn, v proto.Message) error {
 	return nil
 }
 
+var writeBufPool sync.Pool
+
 func write(ctx context.Context, c *websocket.Conn, v proto.Message) error {
-	b, err := proto.Marshal(v)
+	b, ok := writeBufPool.Get().(*proto.Buffer)
+	if !ok {
+		b = proto.NewBuffer(make([]byte, 0, 4096))
+	}
+	defer func() {
+		b.Reset()
+		writeBufPool.Put(b)
+	}()
+
+	err := b.Marshal(v)
 	if err != nil {
 		return xerrors.Errorf("failed to marshal protobuf: %w", err)
 	}
 
-	return c.Write(ctx, websocket.MessageBinary, b)
+	return c.Write(ctx, websocket.MessageBinary, b.Bytes())
 }