gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/container/sync-queue.go (about)

     1  package container
     2  
     3  import (
     4  	"math"
     5  	"sync/atomic"
     6  	"unsafe"
     7  	_ "unsafe"
     8  )
     9  
    10  //go:linkname sync_runtime_Semacquire sync.runtime_Semacquire
    11  func sync_runtime_Semacquire(addr *uint32)
    12  
    13  //go:linkname sync_runtime_Semrelease sync.runtime_Semrelease
    14  func sync_runtime_Semrelease(addr *uint32, handoff bool, skipframes int)
    15  
    16  type SyncQueue[E any] struct {
    17  	raw      unsafe.Pointer
    18  	cap      int64
    19  	mask     int64
    20  	len      atomic.Int64
    21  	head     atomic.Int64
    22  	tail     atomic.Int64
    23  	pushSema uint32
    24  	popSema  uint32
    25  }
    26  
    27  func NewSyncQueue[E any](cap int) *SyncQueue[E] {
    28  	if cap > math.MaxUint32 {
    29  		panic("")
    30  	}
    31  	q := &SyncQueue[E]{cap: int64(cap), pushSema: uint32(cap), popSema: uint32(cap)}
    32  	if cap > 0 {
    33  		if (cap & (cap - 1)) == 0 {
    34  			q.mask = int64(cap - 1)
    35  		}
    36  		raw := make([]E, cap)
    37  		q.raw = unsafe.Pointer(&raw[0])
    38  	}
    39  	return q
    40  }
    41  
    42  func (q *SyncQueue[E]) TryPush(e E) bool {
    43  	if l := q.len.Add(1); l > q.cap {
    44  		q.len.Add(-1)
    45  		return false
    46  	}
    47  	tail := q.tail.Add(1)
    48  	if q.mask > 0 {
    49  		*(*E)(unsafe.Pointer(uintptr(q.raw) + uintptr((tail-1)&q.mask)*unsafe.Sizeof(e))) = e
    50  	} else {
    51  		*(*E)(unsafe.Pointer(uintptr(q.raw) + uintptr((tail-1)%q.cap)*unsafe.Sizeof(e))) = e
    52  	}
    53  	return true
    54  }
    55  
    56  func (q *SyncQueue[E]) TryPop() (e E, ok bool) {
    57  	head := q.head.Add(1)
    58  	if head > q.tail.Load() {
    59  		q.head.Add(-1)
    60  		return
    61  	}
    62  	q.len.Add(-1)
    63  	if q.mask > 0 {
    64  		return *(*E)(unsafe.Pointer(uintptr(q.raw) + uintptr((head-1)&q.mask)*unsafe.Sizeof(e))), true
    65  	} else {
    66  		return *(*E)(unsafe.Pointer(uintptr(q.raw) + uintptr((head-1)%q.cap)*unsafe.Sizeof(e))), true
    67  	}
    68  }