Improve: alloc using make if alloc size > 65536 (#2796)

This commit is contained in:
Kr328 2023-06-18 11:19:35 +08:00 committed by GitHub
parent 295b0da0e5
commit 154cb1d1f0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 28 additions and 14 deletions

View File

@ -32,23 +32,32 @@ func NewAllocator() *Allocator {
// Get a []byte from pool with most appropriate cap // Get a []byte from pool with most appropriate cap
func (alloc *Allocator) Get(size int) []byte { func (alloc *Allocator) Get(size int) []byte {
if size <= 0 || size > 65536 { switch {
case size < 0:
panic("alloc.Get: len out of range")
case size == 0:
return nil return nil
} case size > 65536:
return make([]byte, size)
default:
bits := msb(size)
if size == 1<<bits {
return alloc.buffers[bits].Get().([]byte)[:size]
}
bits := msb(size) return alloc.buffers[bits+1].Get().([]byte)[:size]
if size == 1<<bits {
return alloc.buffers[bits].Get().([]byte)[:size]
} }
return alloc.buffers[bits+1].Get().([]byte)[:size]
} }
// Put returns a []byte to pool for future use, // Put returns a []byte to pool for future use,
// which the cap must be exactly 2^n // which the cap must be exactly 2^n
func (alloc *Allocator) Put(buf []byte) error { func (alloc *Allocator) Put(buf []byte) error {
if cap(buf) == 0 || cap(buf) > 65536 {
return nil
}
bits := msb(cap(buf)) bits := msb(cap(buf))
if cap(buf) == 0 || cap(buf) > 65536 || cap(buf) != 1<<bits { if cap(buf) != 1<<bits {
return errors.New("allocator Put() incorrect buffer size") return errors.New("allocator Put() incorrect buffer size")
} }

View File

@ -19,17 +19,17 @@ func TestAllocGet(t *testing.T) {
assert.Equal(t, 1024, cap(alloc.Get(1023))) assert.Equal(t, 1024, cap(alloc.Get(1023)))
assert.Equal(t, 1024, len(alloc.Get(1024))) assert.Equal(t, 1024, len(alloc.Get(1024)))
assert.Equal(t, 65536, len(alloc.Get(65536))) assert.Equal(t, 65536, len(alloc.Get(65536)))
assert.Nil(t, alloc.Get(65537)) assert.Equal(t, 65537, len(alloc.Get(65537)))
} }
func TestAllocPut(t *testing.T) { func TestAllocPut(t *testing.T) {
alloc := NewAllocator() alloc := NewAllocator()
assert.NotNil(t, alloc.Put(nil), "put nil misbehavior") assert.Nil(t, alloc.Put(nil), "put nil misbehavior")
assert.NotNil(t, alloc.Put(make([]byte, 3)), "put elem:3 []bytes misbehavior") assert.NotNil(t, alloc.Put(make([]byte, 3)), "put elem:3 []bytes misbehavior")
assert.Nil(t, alloc.Put(make([]byte, 4)), "put elem:4 []bytes misbehavior") assert.Nil(t, alloc.Put(make([]byte, 4)), "put elem:4 []bytes misbehavior")
assert.Nil(t, alloc.Put(make([]byte, 1023, 1024)), "put elem:1024 []bytes misbehavior") assert.Nil(t, alloc.Put(make([]byte, 1023, 1024)), "put elem:1024 []bytes misbehavior")
assert.Nil(t, alloc.Put(make([]byte, 65536)), "put elem:65536 []bytes misbehavior") assert.Nil(t, alloc.Put(make([]byte, 65536)), "put elem:65536 []bytes misbehavior")
assert.NotNil(t, alloc.Put(make([]byte, 65537)), "put elem:65537 []bytes misbehavior") assert.Nil(t, alloc.Put(make([]byte, 65537)), "put elem:65537 []bytes misbehavior")
} }
func TestAllocPutThenGet(t *testing.T) { func TestAllocPutThenGet(t *testing.T) {

View File

@ -7,6 +7,8 @@ import (
"unsafe" "unsafe"
"golang.org/x/sys/windows" "golang.org/x/sys/windows"
"github.com/Dreamacro/clash/common/pool"
) )
var ( var (
@ -41,8 +43,10 @@ func findProcessPath(network string, from netip.AddrPort, to netip.AddrPort) (st
} }
func findPidByConnectionEndpoint(family uint32, protocol uint32, from netip.AddrPort, to netip.AddrPort) (uint32, error) { func findPidByConnectionEndpoint(family uint32, protocol uint32, from netip.AddrPort, to netip.AddrPort) (uint32, error) {
buf := []byte(nil) buf := pool.Get(0)
bufSize := uint32(0) defer pool.Put(buf)
bufSize := uint32(len(buf))
loop: loop:
for { for {
@ -77,7 +81,8 @@ loop:
break loop break loop
case uintptr(windows.ERROR_INSUFFICIENT_BUFFER): case uintptr(windows.ERROR_INSUFFICIENT_BUFFER):
buf = make([]byte, bufSize) pool.Put(buf)
buf = pool.Get(int(bufSize))
continue loop continue loop
default: default: