github.com/GoWebProd/gip@v0.0.0-20230623090727-b60d41d5d320/pool/pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"sync/atomic"
     7  	"unsafe"
     8  
     9  	"github.com/GoWebProd/gip/rtime"
    10  )
    11  
    12  var allPoolsMu sync.Mutex
    13  
    14  //go:linkname runtime_LoadAcquintptr runtime/internal/atomic.LoadAcquintptr
    15  func runtime_LoadAcquintptr(ptr *uintptr) uintptr
    16  
    17  //go:linkname runtime_StoreReluintptr runtime/internal/atomic.StoreReluintptr
    18  func runtime_StoreReluintptr(ptr *uintptr, val uintptr) uintptr
    19  
    20  type noCopy struct{}
    21  
    22  type Pool[T any] struct {
    23  	noCopy noCopy
    24  
    25  	local     unsafe.Pointer // local fixed-size per-P pool, actual type is [P]poolLocal
    26  	localSize uintptr        // size of the local array
    27  }
    28  
    29  func (p *Pool[T]) pin() (*poolLocal[T], int) {
    30  	pid := rtime.ProcPin()
    31  	// In pinSlow we store to local and then to localSize, here we load in opposite order.
    32  	// Since we've disabled preemption, GC cannot happen in between.
    33  	// Thus here we must observe local at least as large localSize.
    34  	// We can observe a newer/larger local, it is fine (we must observe its zero-initialized-ness).
    35  	s := runtime_LoadAcquintptr(&p.localSize) // load-acquire
    36  	l := p.local                              // load-consume
    37  	if uintptr(pid) < s {
    38  		return indexLocal[T](l, pid), pid
    39  	}
    40  	return p.pinSlow()
    41  }
    42  
    43  func (p *Pool[T]) pinSlow() (*poolLocal[T], int) {
    44  	// Retry under the mutex.
    45  	// Can not lock the mutex while pinned.
    46  	rtime.ProcUnpin()
    47  	allPoolsMu.Lock()
    48  	defer allPoolsMu.Unlock()
    49  
    50  	pid := rtime.ProcPin()
    51  	// poolCleanup won't be called while we are pinned.
    52  	s := p.localSize
    53  	l := p.local
    54  
    55  	if uintptr(pid) < s {
    56  		return indexLocal[T](l, pid), pid
    57  	}
    58  	size := runtime.GOMAXPROCS(0)
    59  	local := make([]poolLocal[T], size)
    60  
    61  	atomic.StorePointer(&p.local, unsafe.Pointer(&local[0])) // store-release
    62  	runtime_StoreReluintptr(&p.localSize, uintptr(size))     // store-release
    63  
    64  	return &local[pid], pid
    65  }
    66  
    67  // Put adds x to the pool.
    68  func (p *Pool[T]) Put(x *T) {
    69  	if x == nil {
    70  		return
    71  	}
    72  
    73  	l, _ := p.pin()
    74  
    75  	if l.private == nil {
    76  		l.private = x
    77  		x = nil
    78  	} else {
    79  		l.shared.pushHead(x)
    80  	}
    81  
    82  	rtime.ProcUnpin()
    83  }
    84  
    85  // Get selects an arbitrary item from the Pool, removes it from the
    86  // Pool, and returns it to the caller.
    87  // Get may choose to ignore the pool and treat it as empty.
    88  // Callers should not assume any relation between values passed to Put and
    89  // the values returned by Get.
    90  //
    91  // If Get would otherwise return nil and p.New is non-nil, Get returns
    92  // the result of calling p.New.
    93  func (p *Pool[T]) Get() *T {
    94  	l, pid := p.pin()
    95  	x := l.private
    96  	l.private = nil
    97  
    98  	if x == nil {
    99  		// Try to pop the head of the local shard. We prefer
   100  		// the head over the tail for temporal locality of
   101  		// reuse.
   102  		x, _ = l.shared.popHead()
   103  		if x == nil {
   104  			x = p.getSlow(pid)
   105  		}
   106  	}
   107  
   108  	rtime.ProcUnpin()
   109  
   110  	return x
   111  }
   112  
   113  func (p *Pool[T]) getSlow(pid int) *T {
   114  	// See the comment in pin regarding ordering of the loads.
   115  	size := runtime_LoadAcquintptr(&p.localSize) // load-acquire
   116  	locals := p.local                            // load-consume
   117  
   118  	// Try to steal one element from other procs.
   119  	for i := 0; i < int(size); i++ {
   120  		l := indexLocal[T](locals, (pid+i+1)%int(size))
   121  
   122  		if x, _ := l.shared.popTail(); x != nil {
   123  			return x
   124  		}
   125  	}
   126  
   127  	return nil
   128  }