forked from p4gefau1t/trojan-go
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
faireal
committed
Sep 11, 2023
1 parent
d722777
commit 51fb05f
Showing
10 changed files
with
322 additions
and
147 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,64 @@ | ||
| package pool | ||
|
|
||
| // Inspired by https://github.com/xtaci/smux/blob/master/alloc.go | ||
|
|
||
| import ( | ||
| "errors" | ||
| "math/bits" | ||
| "sync" | ||
| ) | ||
|
|
||
| var defaultAllocator = NewAllocator() | ||
|
|
||
| // Allocator for incoming frames, optimized to prevent overwriting after zeroing | ||
| type Allocator struct { | ||
| buffers []sync.Pool | ||
| } | ||
|
|
||
| // NewAllocator initiates a []byte allocator for frames less than 65536 bytes, | ||
| // the waste(memory fragmentation) of space allocation is guaranteed to be | ||
| // no more than 50%. | ||
| func NewAllocator() *Allocator { | ||
| alloc := new(Allocator) | ||
| alloc.buffers = make([]sync.Pool, 17) // 1B -> 64K | ||
| for k := range alloc.buffers { | ||
| i := k | ||
| alloc.buffers[k].New = func() any { | ||
| return make([]byte, 1<<uint32(i)) | ||
| } | ||
| } | ||
| return alloc | ||
| } | ||
|
|
||
| // Get a []byte from pool with most appropriate cap | ||
| func (alloc *Allocator) Get(size int) []byte { | ||
| if size <= 0 || size > 65536 { | ||
| return nil | ||
| } | ||
|
|
||
| bits := msb(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, | ||
| // which the cap must be exactly 2^n | ||
| func (alloc *Allocator) Put(buf []byte) error { | ||
| bits := msb(cap(buf)) | ||
| if cap(buf) == 0 || cap(buf) > 65536 || cap(buf) != 1<<bits { | ||
| return errors.New("allocator Put() incorrect buffer size") | ||
| } | ||
|
|
||
| //nolint | ||
| //lint:ignore SA6002 ignore temporarily | ||
| alloc.buffers[bits].Put(buf) | ||
| return nil | ||
| } | ||
|
|
||
| // msb return the pos of most significant bit | ||
| func msb(size int) uint16 { | ||
| return uint16(bits.Len32(uint32(size)) - 1) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| package pool | ||
|
|
||
| import ( | ||
| "math/rand" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestAllocGet(t *testing.T) { | ||
| alloc := NewAllocator() | ||
| assert.Nil(t, alloc.Get(0)) | ||
| assert.Equal(t, 1, len(alloc.Get(1))) | ||
| assert.Equal(t, 2, len(alloc.Get(2))) | ||
| assert.Equal(t, 3, len(alloc.Get(3))) | ||
| assert.Equal(t, 4, cap(alloc.Get(3))) | ||
| assert.Equal(t, 4, cap(alloc.Get(4))) | ||
| assert.Equal(t, 1023, len(alloc.Get(1023))) | ||
| assert.Equal(t, 1024, cap(alloc.Get(1023))) | ||
| assert.Equal(t, 1024, len(alloc.Get(1024))) | ||
| assert.Equal(t, 65536, len(alloc.Get(65536))) | ||
| assert.Nil(t, alloc.Get(65537)) | ||
| } | ||
|
|
||
| func TestAllocPut(t *testing.T) { | ||
| alloc := NewAllocator() | ||
| assert.NotNil(t, alloc.Put(nil), "put nil 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, 1023, 1024)), "put elem:1024 []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") | ||
| } | ||
|
|
||
| func TestAllocPutThenGet(t *testing.T) { | ||
| alloc := NewAllocator() | ||
| data := alloc.Get(4) | ||
| alloc.Put(data) | ||
| newData := alloc.Get(4) | ||
|
|
||
| assert.Equal(t, cap(data), cap(newData), "different cap while alloc.Get()") | ||
| } | ||
|
|
||
| func BenchmarkMSB(b *testing.B) { | ||
| for i := 0; i < b.N; i++ { | ||
| msb(rand.Int()) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| package pool | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "sync" | ||
|
|
||
| "github.com/Dreamacro/protobytes" | ||
| ) | ||
|
|
||
| var ( | ||
| bufferPool = sync.Pool{New: func() any { return &bytes.Buffer{} }} | ||
| bytesBufferPool = sync.Pool{New: func() any { return &protobytes.BytesWriter{} }} | ||
| ) | ||
|
|
||
| func GetBuffer() *bytes.Buffer { | ||
| return bufferPool.Get().(*bytes.Buffer) | ||
| } | ||
|
|
||
| func PutBuffer(buf *bytes.Buffer) { | ||
| buf.Reset() | ||
| bufferPool.Put(buf) | ||
| } | ||
|
|
||
| func GetBytesBuffer() *protobytes.BytesWriter { | ||
| return bytesBufferPool.Get().(*protobytes.BytesWriter) | ||
| } | ||
|
|
||
| func PutBytesBuffer(buf *protobytes.BytesWriter) { | ||
| buf.Reset() | ||
| bytesBufferPool.Put(buf) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| package pool | ||
|
|
||
| const ( | ||
| // io.Copy default buffer size is 32 KiB | ||
| // but the maximum packet size of vmess/shadowsocks is about 16 KiB | ||
| // so define a buffer of 20 KiB to reduce the memory of each TCP relay | ||
| RelayBufferSize = 20 * 1024 | ||
|
|
||
| // RelayBufferSize uses 20KiB, but due to the allocator it will actually | ||
| // request 32Kib. Most UDPs are smaller than the MTU, and the TUN's MTU | ||
| // set to 9000, so the UDP Buffer size set to 16Kib | ||
| UDPBufferSize = 16 * 1024 | ||
| ) | ||
|
|
||
| func Get(size int) []byte { | ||
| return defaultAllocator.Get(size) | ||
| } | ||
|
|
||
| func Put(buf []byte) error { | ||
| return defaultAllocator.Put(buf) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.