diff --git a/service/rpc/io_stream.go b/service/rpc/io_stream.go index 82a63e7..b1014cb 100644 --- a/service/rpc/io_stream.go +++ b/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 }