gtsocial-umbx

pool.go (8684B)

---

 package pools
 
 import (
 	"runtime"
 	"sync"
 	"sync/atomic"
 	"unsafe"
 )
 
 type Pool struct {
 	// New is used to instantiate new items
 	New func() interface{}
 
 	// Evict is called on evicted items during pool .Clean()
 	Evict func(interface{})
 
 	local    unsafe.Pointer // ptr to []_ppool
 	localSz  int64          // count of all elems in local
 	victim   unsafe.Pointer // ptr to []_ppool
 	victimSz int64          // count of all elems in victim
 	mutex    sync.Mutex     // mutex protects new cleanups, and new allocations of local
 }
 
 // Get attempts to fetch an item from the pool, failing that allocates with supplied .New() function
 func (p *Pool) Get() interface{} {
 	// Get local pool for proc
 	// (also pins proc)
 	pool, pid := p.pin()
 
 	if v := pool.getPrivate(); v != nil {
 		// local _ppool private elem acquired
 		runtime_procUnpin()
 		atomic.AddInt64(&p.localSz, -1)
 		return v
 	}
 
 	if v := pool.get(); v != nil {
 		// local _ppool queue elem acquired
 		runtime_procUnpin()
 		atomic.AddInt64(&p.localSz, -1)
 		return v
 	}
 
 	// Unpin before attempting slow
 	runtime_procUnpin()
 	if v := p.getSlow(pid); v != nil {
 		// note size decrementing
 		// is handled within p.getSlow()
 		// as we don't know if it came
 		// from the local or victim pools
 		return v
 	}
 
 	// Alloc new
 	return p.New()
 }
 
 // Put places supplied item in the proc local pool
 func (p *Pool) Put(v interface{}) {
 	// Don't store nil
 	if v == nil {
 		return
 	}
 
 	// Get proc local pool
 	// (also pins proc)
 	pool, _ := p.pin()
 
 	// first try private, then queue
 	if !pool.setPrivate(v) {
 		pool.put(v)
 	}
 	runtime_procUnpin()
 
 	// Increment local pool size
 	atomic.AddInt64(&p.localSz, 1)
 }
 
 // Clean will drop the current victim pools, move the current local pools to its
 // place and reset the local pools ptr in order to be regenerated
 func (p *Pool) Clean() {
 	p.mutex.Lock()
 
 	// victim becomes local, local becomes nil
 	localPtr := atomic.SwapPointer(&p.local, nil)
 	victimPtr := atomic.SwapPointer(&p.victim, localPtr)
 	localSz := atomic.SwapInt64(&p.localSz, 0)
 	atomic.StoreInt64(&p.victimSz, localSz)
 
 	var victim []ppool
 	if victimPtr != nil {
 		victim = *(*[]ppool)(victimPtr)
 	}
 
 	// drain each of the vict _ppool items
 	for i := 0; i < len(victim); i++ {
 		ppool := &victim[i]
 		ppool.evict(p.Evict)
 	}
 
 	p.mutex.Unlock()
 }
 
 // LocalSize returns the total number of elements in all the proc-local pools
 func (p *Pool) LocalSize() int64 {
 	return atomic.LoadInt64(&p.localSz)
 }
 
 // VictimSize returns the total number of elements in all the victim (old proc-local) pools
 func (p *Pool) VictimSize() int64 {
 	return atomic.LoadInt64(&p.victimSz)
 }
 
 // getSlow is the slow path for fetching an element, attempting to steal from other proc's
 // local pools, and failing that, from the aging-out victim pools. pid is still passed so
 // not all procs start iterating from the same index
 func (p *Pool) getSlow(pid int) interface{} {
 	// get local pools
 	local := p.localPools()
 
 	// Try to steal from other proc locals
 	for i := 0; i < len(local); i++ {
 		pool := &local[(pid+i+1)%len(local)]
 		if v := pool.get(); v != nil {
 			atomic.AddInt64(&p.localSz, -1)
 			return v
 		}
 	}
 
 	// get victim pools
 	victim := p.victimPools()
 
 	// Attempt to steal from victim pools
 	for i := 0; i < len(victim); i++ {
 		pool := &victim[(pid+i+1)%len(victim)]
 		if v := pool.get(); v != nil {
 			atomic.AddInt64(&p.victimSz, -1)
 			return v
 		}
 	}
 
 	// Set victim pools to nil (none found)
 	atomic.StorePointer(&p.victim, nil)
 
 	return nil
 }
 
 // localPools safely loads slice of local _ppools
 func (p *Pool) localPools() []ppool {
 	local := atomic.LoadPointer(&p.local)
 	if local == nil {
 		return nil
 	}
 	return *(*[]ppool)(local)
 }
 
 // victimPools safely loads slice of victim _ppools
 func (p *Pool) victimPools() []ppool {
 	victim := atomic.LoadPointer(&p.victim)
 	if victim == nil {
 		return nil
 	}
 	return *(*[]ppool)(victim)
 }
 
 // pin will get fetch pin proc to PID, fetch proc-local _ppool and current PID we're pinned to
 func (p *Pool) pin() (*ppool, int) {
 	for {
 		// get local pools
 		local := p.localPools()
 
 		if len(local) > 0 {
 			// local already initialized
 
 			// pin to current proc
 			pid := runtime_procPin()
 
 			// check for pid local pool
 			if pid < len(local) {
 				return &local[pid], pid
 			}
 
 			// unpin from proc
 			runtime_procUnpin()
 		} else {
 			// local not yet initialized
 
 			// Check functions are set
 			if p.New == nil {
 				panic("new func must not be nil")
 			}
 			if p.Evict == nil {
 				panic("evict func must not be nil")
 			}
 		}
 
 		// allocate local
 		p.allocLocal()
 	}
 }
 
 // allocLocal allocates a new local pool slice, with the old length passed to check
 // if pool was previously nil, or whether a change in GOMAXPROCS occurred
 func (p *Pool) allocLocal() {
 	// get pool lock
 	p.mutex.Lock()
 
 	// Calculate new size to use
 	size := runtime.GOMAXPROCS(0)
 
 	local := p.localPools()
 	if len(local) != size {
 		// GOMAXPROCS changed, reallocate
 		pools := make([]ppool, size)
 		atomic.StorePointer(&p.local, unsafe.Pointer(&pools))
 
 		// Evict old local elements
 		for i := 0; i < len(local); i++ {
 			pool := &local[i]
 			pool.evict(p.Evict)
 		}
 	}
 
 	// Unlock pool
 	p.mutex.Unlock()
 }
 
 // _ppool is a proc local pool
 type _ppool struct {
 	// root is the root element of the _ppool queue,
 	// and protects concurrent access to the queue
 	root unsafe.Pointer
 
 	// private is a proc private member accessible
 	// only to the pid this _ppool is assigned to,
 	// except during evict (hence the unsafe pointer)
 	private unsafe.Pointer
 }
 
 // ppool wraps _ppool with pad.
 type ppool struct {
 	_ppool
 
 	// Prevents false sharing on widespread platforms with
 	// 128 mod (cache line size) = 0 .
 	pad [128 - unsafe.Sizeof(_ppool{})%128]byte
 }
 
 // getPrivate gets the proc private member
 func (pp *_ppool) getPrivate() interface{} {
 	ptr := atomic.SwapPointer(&pp.private, nil)
 	if ptr == nil {
 		return nil
 	}
 	return *(*interface{})(ptr)
 }
 
 // setPrivate sets the proc private member (only if unset)
 func (pp *_ppool) setPrivate(v interface{}) bool {
 	return atomic.CompareAndSwapPointer(&pp.private, nil, unsafe.Pointer(&v))
 }
 
 // get fetches an element from the queue
 func (pp *_ppool) get() interface{} {
 	for {
 		// Attempt to load root elem
 		root := atomic.LoadPointer(&pp.root)
 		if root == nil {
 			return nil
 		}
 
 		// Attempt to consume root elem
 		if root == inUsePtr ||
 			!atomic.CompareAndSwapPointer(&pp.root, root, inUsePtr) {
 			continue
 		}
 
 		// Root becomes next in chain
 		e := (*elem)(root)
 		v := e.value
 
 		// Place new root back in the chain
 		atomic.StorePointer(&pp.root, unsafe.Pointer(e.next))
 		putElem(e)
 
 		return v
 	}
 }
 
 // put places an element in the queue
 func (pp *_ppool) put(v interface{}) {
 	// Prepare next elem
 	e := getElem()
 	e.value = v
 
 	for {
 		// Attempt to load root elem
 		root := atomic.LoadPointer(&pp.root)
 		if root == inUsePtr {
 			continue
 		}
 
 		// Set the next elem value (might be nil)
 		e.next = (*elem)(root)
 
 		// Attempt to store this new value at root
 		if atomic.CompareAndSwapPointer(&pp.root, root, unsafe.Pointer(e)) {
 			break
 		}
 	}
 }
 
 // hook evicts all entries from pool, calling hook on each
 func (pp *_ppool) evict(hook func(interface{})) {
 	if v := pp.getPrivate(); v != nil {
 		hook(v)
 	}
 	for {
 		v := pp.get()
 		if v == nil {
 			break
 		}
 		hook(v)
 	}
 }
 
 // inUsePtr is a ptr used to indicate _ppool is in use
 var inUsePtr = unsafe.Pointer(&elem{
 	next:  nil,
 	value: "in_use",
 })
 
 // elem defines an element in the _ppool queue
 type elem struct {
 	next  *elem
 	value interface{}
 }
 
 // elemPool is a simple pool of unused elements
 var elemPool = struct {
 	root unsafe.Pointer
 }{}
 
 // getElem fetches a new elem from pool, or creates new
 func getElem() *elem {
 	// Attempt to load root elem
 	root := atomic.LoadPointer(&elemPool.root)
 	if root == nil {
 		return &elem{}
 	}
 
 	// Attempt to consume root elem
 	if root == inUsePtr ||
 		!atomic.CompareAndSwapPointer(&elemPool.root, root, inUsePtr) {
 		return &elem{}
 	}
 
 	// Root becomes next in chain
 	e := (*elem)(root)
 	atomic.StorePointer(&elemPool.root, unsafe.Pointer(e.next))
 	e.next = nil
 
 	return e
 }
 
 // putElem will place element in the pool
 func putElem(e *elem) {
 	e.value = nil
 
 	// Attempt to load root elem
 	root := atomic.LoadPointer(&elemPool.root)
 	if root == inUsePtr {
 		return // drop
 	}
 
 	// Set the next elem value (might be nil)
 	e.next = (*elem)(root)
 
 	// Attempt to store this new value at root
 	atomic.CompareAndSwapPointer(&elemPool.root, root, unsafe.Pointer(e))
 }
 
 //go:linkname runtime_procPin sync.runtime_procPin
 func runtime_procPin() int
 
 //go:linkname runtime_procUnpin sync.runtime_procUnpin
 func runtime_procUnpin()