Use internal pool

dialer.go

@@ -15,7 +15,7 @@ import (
 	"time"
 
 	"github.com/gobwas/httphead"
-	"github.com/gobwas/pool/pbufio"
+	"github.com/gobwas/ws/internal/pbufio"
 )
 
 // Constants used by Dialer.

go.mod

@@ -4,6 +4,5 @@ go 1.15
 
 require (
 	github.com/gobwas/httphead v0.1.0
-	github.com/gobwas/pool v0.2.1
-	golang.org/x/sys v0.6.0 // indirect
+	golang.org/x/sys v0.8.0
 )

go.sum

@@ -1,8 +1,4 @@
 github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU=
 github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
-github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
-github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
-golang.org/x/sys v0.0.0-20201207223542-d4d67f95c62d h1:MiWWjyhUzZ+jvhZvloX6ZrUsdEghn8a64Upd8EMHglE=
-golang.org/x/sys v0.0.0-20201207223542-d4d67f95c62d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
-golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/generic.go

@@ -0,0 +1,161 @@
+package pool
+
+import (
+	"sync"
+)
+
+var defaultPool = New(128, 65536)
+
+// Get pulls object whose generic size is at least of given size. It also
+// returns a real size of x for further pass to Put(). It returns -1 as real
+// size for nil x. Size >-1 does not mean that x is non-nil, so checks must be
+// done.
+//
+// Note that size could be ceiled to the next power of two.
+//
+// Get is a wrapper around defaultPool.Get().
+func Get(size int) (interface{}, int) { return defaultPool.Get(size) }
+
+// Put takes x and its size for future reuse.
+// Put is a wrapper around defaultPool.Put().
+func Put(x interface{}, size int) { defaultPool.Put(x, size) }
+
+// Pool contains logic of reusing objects distinguishable by size in generic
+// way.
+type Pool struct {
+	pool map[int]*sync.Pool
+	size func(int) int
+}
+
+// New creates new Pool that reuses objects which size is in logarithmic range
+// [min, max].
+//
+// Note that it is a shortcut for Custom() constructor with Options provided by
+// WithLogSizeMapping() and WithLogSizeRange(min, max) calls.
+func New(min, max int) *Pool {
+	return Custom(
+		WithLogSizeMapping(),
+		WithLogSizeRange(min, max),
+	)
+}
+
+// Custom creates new Pool with given options.
+func Custom(opts ...Option) *Pool {
+	p := &Pool{
+		pool: make(map[int]*sync.Pool),
+		size: identity,
+	}
+
+	c := (*poolConfig)(p)
+	for _, opt := range opts {
+		opt(c)
+	}
+
+	return p
+}
+
+// Get pulls object whose generic size is at least of given size.
+// It also returns a real size of x for further pass to Put() even if x is nil.
+// Note that size could be ceiled to the next power of two.
+func (p *Pool) Get(size int) (interface{}, int) {
+	n := p.size(size)
+	if pool := p.pool[n]; pool != nil {
+		return pool.Get(), n
+	}
+	return nil, size
+}
+
+// Put takes x and its size for future reuse.
+func (p *Pool) Put(x interface{}, size int) {
+	if pool := p.pool[size]; pool != nil {
+		pool.Put(x)
+	}
+}
+
+type poolConfig Pool
+
+// AddSize adds size n to the map.
+func (p *poolConfig) AddSize(n int) {
+	p.pool[n] = new(sync.Pool)
+}
+
+// SetSizeMapping sets up incoming size mapping function.
+func (p *poolConfig) SetSizeMapping(size func(int) int) {
+	p.size = size
+}
+
+// Option configures pool.
+type Option func(Config)
+
+// Config describes generic pool configuration.
+type Config interface {
+	AddSize(n int)
+	SetSizeMapping(func(int) int)
+}
+
+// WithSizeLogRange returns an Option that will add logarithmic range of
+// pooling sizes containing [min, max] values.
+func WithLogSizeRange(min, max int) Option {
+	return func(c Config) {
+		logarithmicRange(min, max, func(n int) {
+			c.AddSize(n)
+		})
+	}
+}
+
+func WithSizeMapping(sz func(int) int) Option {
+	return func(c Config) {
+		c.SetSizeMapping(sz)
+	}
+}
+
+func WithLogSizeMapping() Option {
+	return WithSizeMapping(ceilToPowerOfTwo)
+}
+
+const (
+	bitsize       = 32 << (^uint(0) >> 63)
+	maxint        = int(1<<(bitsize-1) - 1)
+	maxintHeadBit = 1 << (bitsize - 2)
+)
+
+// logarithmicRange iterates from ceiled to power of two min to max,
+// calling cb on each iteration.
+func logarithmicRange(min, max int, cb func(int)) {
+	if min == 0 {
+		min = 1
+	}
+	for n := ceilToPowerOfTwo(min); n <= max; n <<= 1 {
+		cb(n)
+	}
+}
+
+// identity is identity.
+func identity(n int) int {
+	return n
+}
+
+// ceilToPowerOfTwo returns the least power of two integer value greater than
+// or equal to n.
+func ceilToPowerOfTwo(n int) int {
+	if n&maxintHeadBit != 0 && n > maxintHeadBit {
+		panic("argument is too large")
+	}
+	if n <= 2 {
+		return n
+	}
+	n--
+	n = fillBits(n)
+	n++
+	return n
+}
+
+func fillBits(n int) int {
+	n |= n >> 1
+	n |= n >> 2
+	n |= n >> 4
+	n |= n >> 8
+	n |= n >> 16
+	n |= n >> 32
+	return n
+}

internal/generic_test.go

@@ -0,0 +1,123 @@
+package pool
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+func TestGenericPoolGet(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		min, max int
+		get      int
+		expSize  int
+	}{
+		{
+			min:     0,
+			max:     1,
+			get:     10,
+			expSize: 10,
+		},
+		{
+			min:     0,
+			max:     16,
+			get:     10,
+			expSize: 16,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			p := New(test.min, test.max)
+			_, n := p.Get(test.get)
+			if n != test.expSize {
+				t.Errorf("Get(%d) = _, %d; want %d", test.get, n, test.expSize)
+			}
+		})
+	}
+}
+
+func TestLogarithmicRange(t *testing.T) {
+	for _, test := range []struct {
+		min, max int
+		exp      []int
+	}{
+		{0, 8, []int{1, 2, 4, 8}},
+		{0, 7, []int{1, 2, 4}},
+		{0, 9, []int{1, 2, 4, 8}},
+		{3, 8, []int{4, 8}},
+		{1, 7, []int{1, 2, 4}},
+		{1, 9, []int{1, 2, 4, 8}},
+	} {
+		t.Run("", func(t *testing.T) {
+			var act []int
+			logarithmicRange(test.min, test.max, func(n int) {
+				act = append(act, n)
+			})
+			if !reflect.DeepEqual(act, test.exp) {
+				t.Errorf("unexpected range from %d to %d: %v; want %v", test.min, test.max, act, test.exp)
+			}
+		})
+	}
+}
+
+func TestCeilToPowerOfTwo(t *testing.T) {
+	for _, test := range []struct {
+		in    int
+		exp   int
+		panic bool
+	}{
+		{in: 0, exp: 0},
+		{in: 1, exp: 1},
+		{in: 2, exp: 2},
+		{in: 3, exp: 4},
+		{in: 4, exp: 4},
+		{in: 9, exp: 16},
+
+		{in: maxintHeadBit - 1, exp: maxintHeadBit},
+		{in: maxintHeadBit + 1, panic: true},
+	} {
+		t.Run(fmt.Sprintf("%d to %d", test.in, test.exp), func(t *testing.T) {
+			defer func() {
+				err := recover()
+				if !test.panic && err != nil {
+					t.Fatalf("panic: %v", err)
+				}
+				if test.panic && err == nil {
+					t.Fatalf("want panic")
+				}
+			}()
+			act := ceilToPowerOfTwo(test.in)
+			if exp := test.exp; act != exp {
+				t.Errorf("CeilToPowerOfTwo(%d) = %d; want %d", test.in, act, exp)
+			}
+		})
+	}
+}
+
+func TestFillBits(t *testing.T) {
+	for _, test := range []struct {
+		in  int
+		exp int
+	}{
+		{0, 0},
+		{1, 1},
+		{btoi("0100"), btoi("0111")},
+		{btoi("0101"), btoi("0111")},
+		{maxintHeadBit, maxint},
+	} {
+		t.Run(fmt.Sprintf("%v", test.in), func(t *testing.T) {
+			act := fillBits(test.in)
+			if exp := test.exp; act != exp {
+				t.Errorf(
+					"fillBits(%064b) = %064b; want %064b",
+					test.in, act, exp,
+				)
+			}
+		})
+	}
+}
+
+func btoi(s string) (n int) {
+	fmt.Sscanf(s, "%b", &n)
+	return n
+}