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