fix(grpc): use mcache to fix memory leak caused by grpc codec buffer …

…to reuse memory incorrectly (#992)
cloudwego · Jun 21, 2023 · 1c4ea3a · 1c4ea3a
1 parent c8fe849
commit 1c4ea3a
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/pkg/remote/trans/gonet/trans_server_test.go b/pkg/remote/trans/gonet/trans_server_test.go
@@ -68,7 +68,7 @@ func TestMain(m *testing.M) {
 // TestCreateListener test trans_server CreateListener success
 func TestCreateListener(t *testing.T) {
 	// tcp init
-	addrStr := "127.0.0.1:9091"
+	addrStr := "127.0.0.1:9093"
 	addr = utils.NewNetAddr("tcp", addrStr)
 
 	// test

diff --git a/pkg/remote/trans/netpoll/trans_server_test.go b/pkg/remote/trans/netpoll/trans_server_test.go
@@ -74,7 +74,7 @@ func TestMain(m *testing.M) {
 // TestCreateListener test trans_server CreateListener success
 func TestCreateListener(t *testing.T) {
 	// tcp init
-	addrStr := "127.0.0.1:9090"
+	addrStr := "127.0.0.1:9091"
 	addr = utils.NewNetAddr("tcp", addrStr)
 
 	// test
@@ -98,7 +98,7 @@ func TestCreateListener(t *testing.T) {
 // TestBootStrap test trans_server BootstrapServer success
 func TestBootStrap(t *testing.T) {
 	// tcp init
-	addrStr := "127.0.0.1:9090"
+	addrStr := "127.0.0.1:9092"
 	addr = utils.NewNetAddr("tcp", addrStr)
 
 	// test

diff --git a/pkg/remote/trans/nphttp2/buffer.go b/pkg/remote/trans/nphttp2/buffer.go
@@ -20,9 +20,16 @@ import (
 	"net"
 	"sync"
 
+	"github.com/bytedance/gopkg/lang/mcache"
+
 	"github.com/cloudwego/kitex/pkg/remote"
 )
 
+const (
+	// min is 4k
+	minMallocSize = 4 * 1024
+)
+
 // GRPCConn implement WriteFrame and ReadFrame
 type GRPCConn interface {
 	net.Conn
@@ -45,11 +52,7 @@ var (
 
 var bufferPool = sync.Pool{
 	New: func() interface{} {
-		return &buffer{
-			rbuf: make([]byte, 0, 4096),
-			wbuf: nil,
-			whdr: nil,
-		}
+		return &buffer{}
 	},
 }
 
@@ -64,7 +67,13 @@ func (b *buffer) growRbuf(n int) {
 	if capacity >= n {
 		return
 	}
-	buf := make([]byte, 0, 2*capacity+n)
+	if n < minMallocSize {
+		n = minMallocSize
+	}
+	buf := mcache.Malloc(0, n)
+	if cap(b.rbuf) > 0 {
+		mcache.Free(b.rbuf)
+	}
 	b.rbuf = buf
 }
 
@@ -95,7 +104,11 @@ func (b *buffer) Flush() (err error) {
 }
 
 func (b *buffer) Release(e error) (err error) {
-	b.rbuf = b.rbuf[:0]
+	if cap(b.rbuf) > 0 {
+		mcache.Free(b.rbuf)
+	}
+	b.conn = nil
+	b.rbuf = nil
 	b.whdr = nil
 	b.wbuf = nil
 	bufferPool.Put(b)