Jason Lyu
2 months ago
committed by
GitHub
GitHub
17 changed files with 243 additions and 96 deletions
@ -0,0 +1,29 @@ |
|||
// Package buffer provides a pool of []byte.
|
|||
package buffer |
|||
|
|||
import ( |
|||
"github.com/xjasonlyu/tun2socks/v2/buffer/allocator" |
|||
) |
|||
|
|||
const ( |
|||
// MaxSegmentSize is the largest possible UDP datagram size.
|
|||
MaxSegmentSize = (1 << 16) - 1 |
|||
|
|||
// RelayBufferSize is the default buffer size for TCP relays.
|
|||
// 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 << 10 |
|||
) |
|||
|
|||
var _allocator = allocator.New() |
|||
|
|||
// Get gets a []byte from default allocator with most appropriate cap.
|
|||
func Get(size int) []byte { |
|||
return _allocator.Get(size) |
|||
} |
|||
|
|||
// Put returns a []byte to default allocator for future use.
|
|||
func Put(buf []byte) error { |
|||
return _allocator.Put(buf) |
|||
} |
|||
@ -1,17 +0,0 @@ |
|||
package pool |
|||
|
|||
import ( |
|||
"bytes" |
|||
"sync" |
|||
) |
|||
|
|||
var _bufferPool = sync.Pool{New: func() any { return &bytes.Buffer{} }} |
|||
|
|||
func GetBuffer() *bytes.Buffer { |
|||
return _bufferPool.Get().(*bytes.Buffer) |
|||
} |
|||
|
|||
func PutBuffer(buf *bytes.Buffer) { |
|||
buf.Reset() |
|||
_bufferPool.Put(buf) |
|||
} |
|||
@ -1,23 +1,38 @@ |
|||
// Package pool provides a pool of []byte.
|
|||
// Package pool provides internal pool utilities.
|
|||
package pool |
|||
|
|||
const ( |
|||
// MaxSegmentSize is the largest possible UDP datagram size.
|
|||
MaxSegmentSize = (1 << 16) - 1 |
|||
|
|||
// RelayBufferSize is a buffer of 20 KiB to reduce the memory
|
|||
// of each TCP relay as io.Copy default buffer size is 32 KiB,
|
|||
// but the maximum packet size of vmess/shadowsocks is about
|
|||
// 16 KiB, so define .
|
|||
RelayBufferSize = 20 << 10 |
|||
import ( |
|||
"sync" |
|||
) |
|||
|
|||
// Get gets a []byte from default allocator with most appropriate cap.
|
|||
func Get(size int) []byte { |
|||
return _allocator.Get(size) |
|||
// A Pool is a generic wrapper around [sync.Pool] to provide strongly-typed
|
|||
// object pooling.
|
|||
//
|
|||
// Note that SA6002 (ref: https://staticcheck.io/docs/checks/#SA6002) will
|
|||
// not be detected, so all internal pool use must take care to only store
|
|||
// pointer types.
|
|||
type Pool[T any] struct { |
|||
pool sync.Pool |
|||
} |
|||
|
|||
// New returns a new [Pool] for T, and will use fn to construct new Ts when
|
|||
// the pool is empty.
|
|||
func New[T any](fn func() T) *Pool[T] { |
|||
return &Pool[T]{ |
|||
pool: sync.Pool{ |
|||
New: func() any { |
|||
return fn() |
|||
}, |
|||
}, |
|||
} |
|||
} |
|||
|
|||
// Get gets a T from the pool, or creates a new one if the pool is empty.
|
|||
func (p *Pool[T]) Get() T { |
|||
return p.pool.Get().(T) |
|||
} |
|||
|
|||
// Put returns a []byte to default allocator for future use.
|
|||
func Put(buf []byte) error { |
|||
return _allocator.Put(buf) |
|||
// Put returns x into the pool.
|
|||
func (p *Pool[T]) Put(x T) { |
|||
p.pool.Put(x) |
|||
} |
|||
|
|||
@ -0,0 +1,85 @@ |
|||
package pool |
|||
|
|||
import ( |
|||
"runtime/debug" |
|||
"sync" |
|||
"testing" |
|||
|
|||
"github.com/stretchr/testify/require" |
|||
) |
|||
|
|||
type pooledValue[T any] struct { |
|||
value T |
|||
} |
|||
|
|||
func TestNew(t *testing.T) { |
|||
// Disable GC to avoid the victim cache during the test.
|
|||
defer debug.SetGCPercent(debug.SetGCPercent(-1)) |
|||
|
|||
p := New(func() *pooledValue[string] { |
|||
return &pooledValue[string]{ |
|||
value: "new", |
|||
} |
|||
}) |
|||
|
|||
// Probabilistically, 75% of sync.Pool.Put calls will succeed when -race
|
|||
// is enabled (see ref below); attempt to make this quasi-deterministic by
|
|||
// brute force (i.e., put significantly more objects in the pool than we
|
|||
// will need for the test) in order to avoid testing without race enabled.
|
|||
//
|
|||
// ref: https://cs.opensource.google/go/go/+/refs/tags/go1.20.2:src/sync/pool.go;l=100-103
|
|||
for i := 0; i < 1_000; i++ { |
|||
p.Put(&pooledValue[string]{ |
|||
value: t.Name(), |
|||
}) |
|||
} |
|||
|
|||
// Ensure that we always get the expected value. Note that this must only
|
|||
// run a fraction of the number of times that Put is called above.
|
|||
for i := 0; i < 10; i++ { |
|||
func() { |
|||
x := p.Get() |
|||
defer p.Put(x) |
|||
require.Equal(t, t.Name(), x.value) |
|||
}() |
|||
} |
|||
|
|||
// Depool all objects that might be in the pool to ensure that it's empty.
|
|||
for i := 0; i < 1_000; i++ { |
|||
p.Get() |
|||
} |
|||
|
|||
// Now that the pool is empty, it should use the value specified in the
|
|||
// underlying sync.Pool.New func.
|
|||
require.Equal(t, "new", p.Get().value) |
|||
} |
|||
|
|||
func TestNew_Race(t *testing.T) { |
|||
p := New(func() *pooledValue[int] { |
|||
return &pooledValue[int]{ |
|||
value: -1, |
|||
} |
|||
}) |
|||
|
|||
var wg sync.WaitGroup |
|||
defer wg.Wait() |
|||
|
|||
// Run a number of goroutines that read and write pool object fields to
|
|||
// tease out races.
|
|||
for i := 0; i < 1_000; i++ { |
|||
i := i |
|||
|
|||
wg.Add(1) |
|||
go func() { |
|||
defer wg.Done() |
|||
|
|||
x := p.Get() |
|||
defer p.Put(x) |
|||
|
|||
// Must both read and write the field.
|
|||
if n := x.value; n >= -1 { |
|||
x.value = i |
|||
} |
|||
}() |
|||
} |
|||
} |
|||
@ -0,0 +1,25 @@ |
|||
package bufferpool |
|||
|
|||
import ( |
|||
"bytes" |
|||
|
|||
"github.com/xjasonlyu/tun2socks/v2/internal/pool" |
|||
) |
|||
|
|||
const _size = 1024 // by default, create 1 KiB buffers
|
|||
|
|||
var _pool = pool.New(func() *bytes.Buffer { |
|||
buf := &bytes.Buffer{} |
|||
buf.Grow(_size) |
|||
return buf |
|||
}) |
|||
|
|||
func Get() *bytes.Buffer { |
|||
buf := _pool.Get() |
|||
buf.Reset() |
|||
return buf |
|||
} |
|||
|
|||
func Put(b *bytes.Buffer) { |
|||
_pool.Put(b) |
|||
} |
|||
Loading…
Reference in new issue