improve: use sync.Pool for buffer allocation (#423)

service/rpc/io_stream.go

@@ -3,6 +3,7 @@ package rpc
 import (
 	"errors"
 	"io"
+	"sync"
 	"sync/atomic"
 	"time"
 )
@@ -14,6 +15,18 @@ type ioStreamContext struct {
 	agentIoConnectCh chan struct{}
 }
 
+type bp struct {
+	buf []byte
+}
+
+var bufPool = sync.Pool{
+	New: func() any {
+		return &bp{
+			buf: make([]byte, 1024*1024),
+		}
+	},
+}
+
 func (s *NezhaHandler) CreateStream(streamId string) {
 	s.ioStreamMutex.Lock()
 	defer s.ioStreamMutex.Unlock()
@@ -117,7 +130,9 @@ LOOP:
 	endCh := make(chan struct{})
 
 	go func() {
-		_, innerErr := io.CopyBuffer(stream.userIo, stream.agentIo, make([]byte, 1048576))
+		bp := bufPool.Get().(*bp)
+		defer bufPool.Put(bp)
+		_, innerErr := io.CopyBuffer(stream.userIo, stream.agentIo, bp.buf)
 		if innerErr != nil {
 			err = innerErr
 		}
@@ -126,7 +141,9 @@ LOOP:
 		}
 	}()
 	go func() {
-		_, innerErr := io.CopyBuffer(stream.agentIo, stream.userIo, make([]byte, 1048576))
+		bp := bufPool.Get().(*bp)
+		defer bufPool.Put(bp)
+		_, innerErr := io.CopyBuffer(stream.agentIo, stream.userIo, bp.buf)
 		if innerErr != nil {
 			err = innerErr
 		}