clash/component/pool/pool.go

115 lines
1.7 KiB
Go

package pool
import (
"context"
"runtime"
"time"
)
type Factory = func(context.Context) (any, error)
type entry struct {
elm any
time time.Time
}
type Option func(*pool)
// WithEvict set the evict callback
func WithEvict(cb func(any)) 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 any, maxSize)
}
}
// Pool is for GC, see New for detail
type Pool struct {
*pool
}
type pool struct {
ch chan any
factory Factory
evict func(any)
maxAge int64
}
func (p *pool) GetContext(ctx context.Context) (any, 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() (any, error) {
return p.GetContext(context.Background())
}
func (p *pool) Put(item any) {
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 any, 10),
factory: factory,
}
for _, option := range options {
option(p)
}
P := &Pool{p}
runtime.SetFinalizer(P, recycle)
return P
}