add func: SetBufferPool

lxzan · lxzan · commit 5ed31e20fdb3 · 2024-07-29T13:23:26.000+08:00
diff --git a/conn_test.go b/conn_test.go
@@ -60,6 +60,7 @@ func (c *webSocketMocker) OnMessage(socket *Conn, message *Message) {
 }
 
 func TestOthers(t *testing.T) {
+	SetBufferPool(NewBufferPool(128, 256*1024))
 	conn, _ := net.Pipe()
 	upgrader := NewUpgrader(new(BuiltinEventHandler), nil)
 	socket := &Conn{
diff --git a/init.go b/init.go
@@ -1,19 +1,23 @@
 package gws
 
-import "github.com/lxzan/gws/internal"
+import (
+	"bytes"
+	"github.com/lxzan/gws/internal"
+)
 
 var (
-	framePadding  = frameHeader{}            // &#24103;&#22836;&#22635;&#20805;&#29289;
-	binaryPool    = new(internal.BufferPool) // &#20869;&#23384;&#27744;
-	defaultLogger = new(stdLogger)           // &#40664;&#35748;&#26085;&#24535;&#24037;&#20855;
+	framePadding  = frameHeader{}                // &#24103;&#22836;&#22635;&#20805;&#29289;
+	binaryPool    = NewBufferPool(128, 256*1024) // &#20869;&#23384;&#27744;
+	defaultLogger = new(stdLogger)               // &#40664;&#35748;&#26085;&#24535;&#24037;&#20855;
 )
 
-func init() {
-	SetBufferPool(256 * 1024)
+type BufferPool interface {
+	Get(n int) *bytes.Buffer
+	Put(b *bytes.Buffer)
 }
 
-// SetBufferPool set up the memory pool, any memory that exceeds maxSize will not be reclaimed.
-// &#35774;&#32622;&#20869;&#23384;&#27744;, &#36229;&#36807;maxSize&#23558;&#19981;&#20250;&#34987;&#22238;&#25910;.
-func SetBufferPool(maxSize uint32) {
-	binaryPool = internal.NewBufferPool(128, maxSize)
+func NewBufferPool(minSize, maxSize uint32) BufferPool {
+	return internal.NewBufferPool(minSize, maxSize)
 }
+
+func SetBufferPool(p BufferPool) { binaryPool = p }

Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,7 @@ func (c webSocketMocker) OnMessage(socket Conn, message *Message) {`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`func TestOthers(t *testing.T) {`
	`63`	`+ SetBufferPool(NewBufferPool(128, 256*1024))`
`63`	`64`	`conn, _ := net.Pipe()`
`64`	`65`	`upgrader := NewUpgrader(new(BuiltinEventHandler), nil)`
`65`	`66`	`socket := &Conn{`