Feature: support snell v2 (#952)

Co-authored-by: Dreamacro <8615343+Dreamacro@users.noreply.github.com>

component/pool/pool.go

@@ -0,0 +1,114 @@
+package pool
+
+import (
+	"context"
+	"runtime"
+	"time"
+)
+
+type Factory = func(context.Context) (interface{}, error)
+
+type entry struct {
+	elm  interface{}
+	time time.Time
+}
+
+type Option func(*pool)
+
+// WithEvict set the evict callback
+func WithEvict(cb func(interface{})) Option {
+	return func(p *pool) {
+		p.evict = cb
+	}
+}
+
+// WithAge defined element max age (millisecond)
+func WithAge(maxAge int64) Option {
+	return func(p *pool) {
+		p.maxAge = maxAge
+	}
+}
+
+// WithSize defined max size of Pool
+func WithSize(maxSize int) Option {
+	return func(p *pool) {
+		p.ch = make(chan interface{}, maxSize)
+	}
+}
+
+// Pool is for GC, see New for detail
+type Pool struct {
+	*pool
+}
+
+type pool struct {
+	ch      chan interface{}
+	factory Factory
+	evict   func(interface{})
+	maxAge  int64
+}
+
+func (p *pool) GetContext(ctx context.Context) (interface{}, error) {
+	now := time.Now()
+	for {
+		select {
+		case item := <-p.ch:
+			elm := item.(*entry)
+			if p.maxAge != 0 && now.Sub(item.(*entry).time).Milliseconds() > p.maxAge {
+				if p.evict != nil {
+					p.evict(elm.elm)
+				}
+				continue
+			}
+
+			return elm.elm, nil
+		default:
+			return p.factory(ctx)
+		}
+	}
+}
+
+func (p *pool) Get() (interface{}, error) {
+	return p.GetContext(context.Background())
+}
+
+func (p *pool) Put(item interface{}) {
+	e := &entry{
+		elm:  item,
+		time: time.Now(),
+	}
+
+	select {
+	case p.ch <- e:
+		return
+	default:
+		// pool is full
+		if p.evict != nil {
+			p.evict(item)
+		}
+		return
+	}
+}
+
+func recycle(p *Pool) {
+	for item := range p.pool.ch {
+		if p.pool.evict != nil {
+			p.pool.evict(item.(*entry).elm)
+		}
+	}
+}
+
+func New(factory Factory, options ...Option) *Pool {
+	p := &pool{
+		ch:      make(chan interface{}, 10),
+		factory: factory,
+	}
+
+	for _, option := range options {
+		option(p)
+	}
+
+	P := &Pool{p}
+	runtime.SetFinalizer(P, recycle)
+	return P
+}

component/pool/pool_test.go

@@ -0,0 +1,96 @@
+package pool
+
+import (
+	"context"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func lg() Factory {
+	initial := -1
+	return func(context.Context) (interface{}, error) {
+		initial++
+		return initial, nil
+	}
+}
+
+func TestPool_Basic(t *testing.T) {
+	g := lg()
+	pool := New(g)
+
+	elm, _ := pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(elm)
+	elm, _ = pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	elm, _ = pool.Get()
+	assert.Equal(t, 1, elm.(int))
+}
+
+func TestPool_MaxSize(t *testing.T) {
+	g := lg()
+	size := 5
+	pool := New(g, WithSize(size))
+
+	items := []interface{}{}
+
+	for i := 0; i < size; i++ {
+		item, _ := pool.Get()
+		items = append(items, item)
+	}
+
+	extra, _ := pool.Get()
+	assert.Equal(t, size, extra.(int))
+
+	for _, item := range items {
+		pool.Put(item)
+	}
+
+	pool.Put(extra)
+
+	for _, item := range items {
+		elm, _ := pool.Get()
+		assert.Equal(t, item.(int), elm.(int))
+	}
+}
+
+func TestPool_MaxAge(t *testing.T) {
+	g := lg()
+	pool := New(g, WithAge(20))
+
+	elm, _ := pool.Get()
+	pool.Put(elm)
+
+	elm, _ = pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(elm)
+
+	time.Sleep(time.Millisecond * 22)
+	elm, _ = pool.Get()
+	assert.Equal(t, 1, elm.(int))
+}
+
+func TestPool_AutoGC(t *testing.T) {
+	g := lg()
+
+	sign := make(chan int)
+	pool := New(g, WithEvict(func(item interface{}) {
+		sign <- item.(int)
+	}))
+
+	elm, _ := pool.Get()
+	assert.Equal(t, 0, elm.(int))
+	pool.Put(2)
+
+	runtime.GC()
+
+	select {
+	case num := <-sign:
+		assert.Equal(t, 2, num)
+	case <-time.After(time.Second * 3):
+		assert.Fail(t, "something wrong")
+	}
+}