github.com/reusee/pr3@v0.0.0-20240520031754-49012a37a83e/pool.go (about)

     1  package pr3
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	_ "unsafe"
     7  )
     8  
     9  type Pool[T any] struct {
    10  	elems    []PoolElem[T]
    11  	capacity uint32
    12  	fallback sync.Pool
    13  }
    14  
    15  type PoolElem[T any] struct {
    16  	from  *sync.Pool
    17  	refs  atomic.Int32
    18  	value T
    19  }
    20  
    21  func (p *PoolElem[T]) Put() bool {
    22  	if c := p.refs.Add(-1); c == 0 {
    23  		if p.from != nil {
    24  			p.from.Put(p)
    25  		}
    26  		return true
    27  	} else if c < 0 {
    28  		panic("bad put")
    29  	}
    30  	return false
    31  }
    32  
    33  func (p *PoolElem[T]) Inc() {
    34  	p.refs.Add(1)
    35  }
    36  
    37  func NewPool[T any](
    38  	capacity uint32,
    39  	newFunc func() T,
    40  ) *Pool[T] {
    41  	pool := &Pool[T]{
    42  		capacity: capacity,
    43  	}
    44  
    45  	pool.fallback = sync.Pool{
    46  		New: func() any {
    47  			elem := &PoolElem[T]{
    48  				value: newFunc(),
    49  				from:  &pool.fallback,
    50  			}
    51  			return elem
    52  		},
    53  	}
    54  
    55  	pool.elems = make([]PoolElem[T], capacity)
    56  	for i := uint32(0); i < capacity; i++ {
    57  		pool.elems[i].value = newFunc()
    58  	}
    59  
    60  	return pool
    61  }
    62  
    63  func (p *Pool[T]) Get(ptr *T) *PoolElem[T] {
    64  	idx := fastrand() % p.capacity
    65  	if p.elems[idx].refs.CompareAndSwap(0, 1) {
    66  		*ptr = p.elems[idx].value
    67  		return &p.elems[idx]
    68  	}
    69  	return p.fallbackGet(ptr)
    70  }
    71  
    72  func (p *Pool[T]) fallbackGet(ptr *T) *PoolElem[T] {
    73  	elem := p.fallback.Get().(*PoolElem[T])
    74  	elem.refs.Store(1)
    75  	*ptr = elem.value
    76  	return elem
    77  }
    78  
    79  //go:linkname fastrand runtime.fastrand
    80  func fastrand() uint32