github.com/wfusion/gofusion@v1.1.14/common/utils/pool.go (about)

     1  package utils
     2  
     3  import (
     4  	"bytes"
     5  	"github.com/spf13/cast"
     6  	"sync"
     7  )
     8  
     9  const (
    10  	maxBytesPoolable = 16*1024 + 512 // 16.5kb
    11  )
    12  
    13  var (
    14  	// BytesBufferPool 64 is bytes.Buffer smallBufferSize, which is an initial allocation minimal capacity.
    15  	BytesBufferPool = NewPool(
    16  		func() *bytes.Buffer { return bytes.NewBuffer(make([]byte, 0, 64)) },
    17  		PoolableEvictFunc(poolBytesBufferEvict),
    18  	)
    19  	BytesPool = NewPool(
    20  		func() poolBytes { return make([]byte, 0, 64) },
    21  		PoolableEvictFunc(poolBytesEvict),
    22  	)
    23  )
    24  
    25  type Poolable[T any] interface {
    26  	Get(initialized any) (T, func())
    27  	Put(obj T)
    28  }
    29  
    30  type poolableOption[T any] struct {
    31  	evict func(obj T) bool
    32  }
    33  
    34  func PoolableEvictFunc[T any](fn func(obj T) bool) OptionFunc[poolableOption[T]] {
    35  	return func(o *poolableOption[T]) {
    36  		o.evict = fn
    37  	}
    38  }
    39  
    40  type poolResettableA[T any] interface{ Reset(obj any) T }
    41  type poolResettableB interface{ Reset() }
    42  type poolResettableC interface{ Reset() error }
    43  type poolResettableD interface{ Reset(obj any) }
    44  type poolResettableE interface{ Reset(obj any) error }
    45  type poolResettableF[T any] interface{ Reset() T }
    46  
    47  func NewPool[T any](newFn func() T, opts ...OptionExtender) Poolable[T] {
    48  	opt := ApplyOptions[poolableOption[T]](opts...)
    49  	return &poolSealer[T]{
    50  		option: opt,
    51  		newFn:  newFn,
    52  		inner: &sync.Pool{
    53  			New: func() any {
    54  				return any(newFn())
    55  			},
    56  		},
    57  	}
    58  }
    59  
    60  type poolSealer[T any] struct {
    61  	option *poolableOption[T]
    62  	inner  *sync.Pool
    63  	newFn  func() T
    64  }
    65  
    66  func (p *poolSealer[T]) Get(initialized any) (T, func()) {
    67  	obj, ok := p.inner.Get().(T)
    68  	if !ok {
    69  		obj = p.newFn()
    70  	}
    71  
    72  	switch resettable := any(obj).(type) {
    73  	case poolResettableA[T]:
    74  		obj = resettable.Reset(initialized)
    75  	case poolResettableB:
    76  		resettable.Reset()
    77  	case poolResettableC:
    78  		MustSuccess(resettable.Reset())
    79  	case poolResettableD:
    80  		resettable.Reset(initialized)
    81  	case poolResettableE:
    82  		MustSuccess(resettable.Reset(initialized))
    83  	case poolResettableF[T]:
    84  		obj = resettable.Reset()
    85  	}
    86  
    87  	once := new(sync.Once)
    88  	return obj, func() {
    89  		once.Do(func() {
    90  			if p.option.evict == nil || !p.option.evict(obj) {
    91  				p.Put(obj)
    92  			}
    93  		})
    94  	}
    95  }
    96  
    97  func (p *poolSealer[T]) Put(obj T) { p.inner.Put(obj) }
    98  
    99  type poolBytes []byte
   100  
   101  func (p poolBytes) Reset(initLen any) poolBytes {
   102  	iLen := cast.ToInt(initLen)
   103  	pp := p
   104  	if cap(p) < cast.ToInt(iLen) {
   105  		pp = make([]byte, iLen)
   106  	}
   107  
   108  	return pp[:iLen]
   109  }
   110  
   111  func poolBytesBufferEvict(b *bytes.Buffer) bool {
   112  	return b.Cap() >= maxBytesPoolable
   113  }
   114  
   115  func poolBytesEvict(p poolBytes) bool {
   116  	return cap(p) >= maxBytesPoolable
   117  }