reduce buffer allocs

lxzan · Aug 7, 2023 · 9f0ea5b · 9f0ea5b
1 parent 7c1bb91
commit 9f0ea5b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/internal/utils.go b/internal/utils.go
@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"bytes"
 	"crypto/sha1"
 	"encoding/base64"
 	"encoding/binary"
@@ -197,3 +198,15 @@ func ToBinaryNumber[T Integer](n T) T {
 	}
 	return x
 }
+
+type Buffer struct {
+	buf      []byte
+	off      int
+	lastRead int8
+}
+
+//go:nosplit
+func ResetBuffer(b *bytes.Buffer, p []byte) {
+	buffer := (*Buffer)(unsafe.Pointer(b))
+	buffer.buf = p
+}
diff --git a/internal/utils_test.go b/internal/utils_test.go
@@ -210,3 +210,11 @@ func TestGetIntersectionElem(t *testing.T) {
 		assert.Equal(t, "", GetIntersectionElem(b, nil))
 	}
 }
+
+func TestResetBuffer(t *testing.T) {
+	var buf = bytes.NewBufferString("hello")
+	var p = buf.Bytes()
+	p = append(p, "world"...)
+	ResetBuffer(buf, p)
+	assert.Equal(t, "helloworld", buf.String())
+}
diff --git a/reader.go b/reader.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
-
 	"github.com/lxzan/gws/internal"
 )
 
@@ -109,7 +108,8 @@ func (c *Conn) readMessage() error {
 	}
 
 	if fin && opcode != OpcodeContinuation {
-		return c.emitMessage(&Message{index: index, Opcode: opcode, Data: bytes.NewBuffer(p)}, compressed)
+		internal.ResetBuffer(buf, p)
+		return c.emitMessage(&Message{index: index, Opcode: opcode, Data: buf}, compressed)
 	}
 
 	if !fin && opcode != OpcodeContinuation {