github.com/chwjbn/xclash@v0.2.0/component/pool/pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"time"
     7  )
     8  
     9  type Factory = func(context.Context) (interface{}, error)
    10  
    11  type entry struct {
    12  	elm  interface{}
    13  	time time.Time
    14  }
    15  
    16  type Option func(*pool)
    17  
    18  // WithEvict set the evict callback
    19  func WithEvict(cb func(interface{})) Option {
    20  	return func(p *pool) {
    21  		p.evict = cb
    22  	}
    23  }
    24  
    25  // WithAge defined element max age (millisecond)
    26  func WithAge(maxAge int64) Option {
    27  	return func(p *pool) {
    28  		p.maxAge = maxAge
    29  	}
    30  }
    31  
    32  // WithSize defined max size of Pool
    33  func WithSize(maxSize int) Option {
    34  	return func(p *pool) {
    35  		p.ch = make(chan interface{}, maxSize)
    36  	}
    37  }
    38  
    39  // Pool is for GC, see New for detail
    40  type Pool struct {
    41  	*pool
    42  }
    43  
    44  type pool struct {
    45  	ch      chan interface{}
    46  	factory Factory
    47  	evict   func(interface{})
    48  	maxAge  int64
    49  }
    50  
    51  func (p *pool) GetContext(ctx context.Context) (interface{}, error) {
    52  	now := time.Now()
    53  	for {
    54  		select {
    55  		case item := <-p.ch:
    56  			elm := item.(*entry)
    57  			if p.maxAge != 0 && now.Sub(item.(*entry).time).Milliseconds() > p.maxAge {
    58  				if p.evict != nil {
    59  					p.evict(elm.elm)
    60  				}
    61  				continue
    62  			}
    63  
    64  			return elm.elm, nil
    65  		default:
    66  			return p.factory(ctx)
    67  		}
    68  	}
    69  }
    70  
    71  func (p *pool) Get() (interface{}, error) {
    72  	return p.GetContext(context.Background())
    73  }
    74  
    75  func (p *pool) Put(item interface{}) {
    76  	e := &entry{
    77  		elm:  item,
    78  		time: time.Now(),
    79  	}
    80  
    81  	select {
    82  	case p.ch <- e:
    83  		return
    84  	default:
    85  		// pool is full
    86  		if p.evict != nil {
    87  			p.evict(item)
    88  		}
    89  		return
    90  	}
    91  }
    92  
    93  func recycle(p *Pool) {
    94  	for item := range p.pool.ch {
    95  		if p.pool.evict != nil {
    96  			p.pool.evict(item.(*entry).elm)
    97  		}
    98  	}
    99  }
   100  
   101  func New(factory Factory, options ...Option) *Pool {
   102  	p := &pool{
   103  		ch:      make(chan interface{}, 10),
   104  		factory: factory,
   105  	}
   106  
   107  	for _, option := range options {
   108  		option(p)
   109  	}
   110  
   111  	P := &Pool{p}
   112  	runtime.SetFinalizer(P, recycle)
   113  	return P
   114  }