feat: use sync.Pool to reduce alloc

ua/buffer.go

@@ -5,6 +5,7 @@
 package ua
 
 import (
+	"bytes"
 	"encoding/binary"
 	"io"
 	"math"
@@ -156,7 +157,7 @@ func (b *Buffer) ReadBytes() []byte {
 	if b.err != nil {
 		return nil
 	}
-	return d
+	return bytes.Clone(d)
 }
 
 func (b *Buffer) ReadStruct(r interface{}) {

uacp/buffer.go

@@ -0,0 +1,34 @@
+package uacp
+
+import "sync"
+
+type Buffer struct {
+	buf []byte
+}
+
+func (b *Buffer) Bytes() []byte {
+	return b.buf
+}
+
+var bufferPool sync.Pool
+
+func AllocBuffer(size int) *Buffer {
+	v := bufferPool.Get()
+	if v == nil {
+		return &Buffer{
+			buf: make([]byte, size),
+		}
+	}
+	buf := v.(*Buffer)
+	if cap(buf.buf) < size {
+		buf.buf = make([]byte, size)
+	} else {
+		buf.buf = buf.buf[:size]
+	}
+	return buf
+}
+
+func FreeBuffer(b *Buffer) {
+	b.buf = b.buf[:0]
+	bufferPool.Put(b)
+}

uacp/conn.go

@@ -237,7 +237,9 @@ func (c *Conn) Handshake(ctx context.Context, endpoint string) error {
 		return err
 	}
 
-	b, err := c.Receive()
+	buf := AllocBuffer(int(c.ack.ReceiveBufSize))
+	defer FreeBuffer(buf)
+	b, err := c.Receive(buf.Bytes())
 	if err != nil {
 		return err
 	}
@@ -282,7 +284,9 @@ func (c *Conn) Handshake(ctx context.Context, endpoint string) error {
 }
 
 func (c *Conn) srvhandshake(endpoint string) error {
-	b, err := c.Receive()
+	buf := AllocBuffer(int(c.ack.ReceiveBufSize))
+	defer FreeBuffer(buf)
+	b, err := c.Receive(buf.Bytes())
 	if err != nil {
 		c.SendError(ua.StatusBadTCPInternalError)
 		return err
@@ -350,11 +354,7 @@ const hdrlen = 8
 // Receive reads a full UACP message from the underlying connection.
 // The size of b must be at least ReceiveBufSize. Otherwise,
 // the function returns an error.
-func (c *Conn) Receive() ([]byte, error) {
-	// TODO(kung-foo): allow user-specified buffer
-	// TODO(kung-foo): sync.Pool
-	b := make([]byte, c.ack.ReceiveBufSize)
-
+func (c *Conn) Receive(b []byte) ([]byte, error) {
 	if _, err := io.ReadFull(c, b[:hdrlen]); err != nil {
 		// todo(fs): do not wrap this error since it hides io.EOF
 		// todo(fs): use golang.org/x/xerrors

uacp/conn_test.go

@@ -114,7 +114,9 @@ NEXT:
 		t.Fatal(err)
 	}
 
-	got, err := srvConn.Receive()
+	buf := AllocBuffer(int(srvConn.ack.ReceiveBufSize))
+	defer FreeBuffer(buf)
+	got, err := srvConn.Receive(buf.Bytes())
 	if err != nil {
 		t.Fatal(err)
 	}

uasc/secure_channel.go

@@ -337,7 +337,9 @@ func (s *SecureChannel) receive(ctx context.Context) *response {
 
 func (s *SecureChannel) readChunk() (*MessageChunk, error) {
 	// read a full message from the underlying conn.
-	b, err := s.c.Receive()
+	buf := uacp.AllocBuffer(int(s.c.ReceiveBufSize()))
+	defer uacp.FreeBuffer(buf)
+	b, err := s.c.Receive(buf.Bytes())
 	if err == io.EOF || len(b) == 0 {
 		return nil, io.EOF
 	}