From 17c71df594f91e16a494870eda25ca27c8319b54 Mon Sep 17 00:00:00 2001 From: yuanliang Date: Thu, 6 Jun 2024 16:14:16 +0800 Subject: [PATCH 1/2] feat: use sync.Pool to reduce alloc Signed-off-by: yuanliang --- uacp/buffer.go | 34 ++++++++++++++++++++++++++++++++++ uacp/conn.go | 14 +++++++------- uacp/conn_test.go | 4 +++- uasc/secure_channel.go | 4 +++- 4 files changed, 47 insertions(+), 9 deletions(-) create mode 100644 uacp/buffer.go diff --git a/uacp/buffer.go b/uacp/buffer.go new file mode 100644 index 00000000..f8224e1b --- /dev/null +++ b/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) +} diff --git a/uacp/conn.go b/uacp/conn.go index c2acb79c..2a4f56ee 100644 --- a/uacp/conn.go +++ 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 diff --git a/uacp/conn_test.go b/uacp/conn_test.go index 2269bbe2..fd57e37e 100644 --- a/uacp/conn_test.go +++ b/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) } diff --git a/uasc/secure_channel.go b/uasc/secure_channel.go index 78822642..90424958 100644 --- a/uasc/secure_channel.go +++ b/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 } From cd497c958843a33db28a5dab2ffbcd9195d2f5a8 Mon Sep 17 00:00:00 2001 From: yuanliang Date: Fri, 21 Jun 2024 23:04:56 +0800 Subject: [PATCH 2/2] fix: intergation test bug Signed-off-by: yuanliang --- ua/buffer.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ua/buffer.go b/ua/buffer.go index d91b0699..4025c91f 100644 --- a/ua/buffer.go +++ b/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{}) {