github.com/kelleygo/clashcore@v1.0.2/component/pool/pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"time"
     7  )
     8  
     9  type Factory[T any] func(context.Context) (T, error)
    10  
    11  type entry[T any] struct {
    12  	elm  T
    13  	time time.Time
    14  }
    15  
    16  type Option[T any] func(*pool[T])
    17  
    18  // WithEvict set the evict callback
    19  func WithEvict[T any](cb func(T)) Option[T] {
    20  	return func(p *pool[T]) {
    21  		p.evict = cb
    22  	}
    23  }
    24  
    25  // WithAge defined element max age (millisecond)
    26  func WithAge[T any](maxAge int64) Option[T] {
    27  	return func(p *pool[T]) {
    28  		p.maxAge = maxAge
    29  	}
    30  }
    31  
    32  // WithSize defined max size of Pool
    33  func WithSize[T any](maxSize int) Option[T] {
    34  	return func(p *pool[T]) {
    35  		p.ch = make(chan *entry[T], maxSize)
    36  	}
    37  }
    38  
    39  // Pool is for GC, see New for detail
    40  type Pool[T any] struct {
    41  	*pool[T]
    42  }
    43  
    44  type pool[T any] struct {
    45  	ch      chan *entry[T]
    46  	factory Factory[T]
    47  	evict   func(T)
    48  	maxAge  int64
    49  }
    50  
    51  func (p *pool[T]) GetContext(ctx context.Context) (T, error) {
    52  	now := time.Now()
    53  	for {
    54  		select {
    55  		case item := <-p.ch:
    56  			elm := item
    57  			if p.maxAge != 0 && now.Sub(item.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[T]) Get() (T, error) {
    72  	return p.GetContext(context.Background())
    73  }
    74  
    75  func (p *pool[T]) Put(item T) {
    76  	e := &entry[T]{
    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[T any](p *Pool[T]) {
    94  	for item := range p.pool.ch {
    95  		if p.pool.evict != nil {
    96  			p.pool.evict(item.elm)
    97  		}
    98  	}
    99  }
   100  
   101  func New[T any](factory Factory[T], options ...Option[T]) *Pool[T] {
   102  	p := &pool[T]{
   103  		ch:      make(chan *entry[T], 10),
   104  		factory: factory,
   105  	}
   106  
   107  	for _, option := range options {
   108  		option(p)
   109  	}
   110  
   111  	P := &Pool[T]{p}
   112  	runtime.SetFinalizer(P, recycle[T])
   113  	return P
   114  }