tcjlj/nezha
1
0
Fork 0
mirror of https://github.com/nezhahq/nezha.git synced 2025-01-22 12:48:14 -05:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
UUBulb 2024-09-27 18:50:19 +08:00 committed by GitHub
parent f32f127dfc
commit bdf36276da
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 19 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
service/rpc/io_stream.go
View File

@ -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
}