reuse bufio.Reader

lxzan · Sep 5, 2023 · de3bd8d · de3bd8d
1 parent 90819f9
commit de3bd8d
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 2 deletions.
diff --git a/conn.go b/conn.go
@@ -56,6 +56,12 @@ func (c *Conn) ReadLoop() {
 		}
 	}
 	c.handler.OnClose(c, c.err.Load().(error))
+
+	// 回收资源
+	if c.isServer {
+		c.config.readerPool.Put(c.br)
+		c.br = nil
+	}
 }
 
 func (c *Conn) isTextValid(opcode Opcode, payload []byte) bool {

diff --git a/internal/pool.go b/internal/pool.go
@@ -58,3 +58,15 @@ func (p *BufferPool) Get(n int) (*bytes.Buffer, int) {
 	}
 	return bytes.NewBuffer(make([]byte, 0, n)), 0
 }
+
+func NewPool[T any](f func() T) *Pool[T] {
+	return &Pool[T]{p: sync.Pool{New: func() any { return f() }}}
+}
+
+type Pool[T any] struct {
+	p sync.Pool
+}
+
+func (c *Pool[T]) Put(v T) { c.p.Put(v) }
+
+func (c *Pool[T]) Get() T { return c.p.Get().(T) }
diff --git a/internal/pool_test.go b/internal/pool_test.go
@@ -53,3 +53,11 @@ func TestBufferPool(t *testing.T) {
 	buffer, _ := pool.Get(256 * 1024)
 	as.GreaterOrEqual(buffer.Cap(), 256*1024)
 }
+
+func TestPool(t *testing.T) {
+	var p = NewPool(func() int {
+		return 0
+	})
+	assert.Equal(t, 0, p.Get())
+	p.Put(1)
+}
diff --git a/option.go b/option.go
@@ -1,6 +1,7 @@
 package gws
 
 import (
+	"bufio"
 	"compress/flate"
 	"crypto/tls"
 	"github.com/lxzan/gws/internal"
@@ -24,6 +25,7 @@ const (
 
 type (
 	Config struct {
+		readerPool    *internal.Pool[*bufio.Reader]
 		compressors   *compressors
 		decompressors *decompressors
 
@@ -158,6 +160,7 @@ func initServerOption(c *ServerOption) *ServerOption {
 	c.CompressorNum = internal.ToBinaryNumber(c.CompressorNum)
 
 	c.config = &Config{
+		readerPool:          internal.NewPool(func() *bufio.Reader { return bufio.NewReaderSize(nil, c.ReadBufferSize) }),
 		ReadAsyncEnabled:    c.ReadAsyncEnabled,
 		ReadAsyncGoLimit:    c.ReadAsyncGoLimit,
 		ReadMaxPayloadSize:  c.ReadMaxPayloadSize,

diff --git a/option_test.go b/option_test.go
@@ -22,6 +22,7 @@ func validateServerOption(as *assert.Assertions, u *Upgrader) {
 	as.Equal(config.ReadBufferSize, option.ReadBufferSize)
 	as.Equal(config.WriteBufferSize, option.WriteBufferSize)
 	as.Equal(config.CompressorNum, option.CompressorNum)
+	as.NotNil(config.readerPool)
 
 	_, ok := u.option.NewSessionStorage().(*sliceMap)
 	as.True(ok)
@@ -39,6 +40,7 @@ func validateClientOption(as *assert.Assertions, option *ClientOption) {
 	as.Equal(config.CheckUtf8Enabled, option.CheckUtf8Enabled)
 	as.Equal(config.ReadBufferSize, option.ReadBufferSize)
 	as.Equal(config.WriteBufferSize, option.WriteBufferSize)
+	as.Nil(config.readerPool)
 
 	_, ok := option.NewSessionStorage().(*sliceMap)
 	as.True(ok)

diff --git a/updrader.go b/updrader.go
@@ -77,7 +77,9 @@ func (c *Upgrader) hijack(w http.ResponseWriter) (net.Conn, *bufio.Reader, error
 	if err != nil {
 		return nil, nil, err
 	}
-	return netConn, bufio.NewReaderSize(netConn, c.option.ReadBufferSize), nil
+	br := c.option.config.readerPool.Get()
+	br.Reset(netConn)
+	return netConn, br, nil
 }
 
 func (c *Upgrader) doUpgrade(r *http.Request, netConn net.Conn, br *bufio.Reader) (*Conn, error) {
@@ -194,7 +196,8 @@ func (c *Server) RunListener(listener net.Listener) error {
 		}
 
 		go func(conn net.Conn) {
-			br := bufio.NewReaderSize(conn, c.upgrader.option.ReadBufferSize)
+			br := c.upgrader.option.config.readerPool.Get()
+			br.Reset(conn)
 			r, err := http.ReadRequest(br)
 			if err != nil {
 				c.OnError(conn, err)