github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/internal/benchmark/msq/msq.go (about)

     1  package msq
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"unsafe"
     7  )
     8  
     9  var msqv1pool *sync.Pool = &sync.Pool{New: func() interface{} { return new(msqv1node) }}
    10  
    11  type MSQueue struct {
    12  	head unsafe.Pointer // *msqv1node
    13  	tail unsafe.Pointer // *msqv1node
    14  }
    15  
    16  type msqv1node struct {
    17  	value uint64
    18  	next  unsafe.Pointer // *msqv1node
    19  }
    20  
    21  func New() *MSQueue {
    22  	node := unsafe.Pointer(new(msqv1node))
    23  	return &MSQueue{head: node, tail: node}
    24  }
    25  
    26  func loadMSQPointer(p *unsafe.Pointer) *msqv1node {
    27  	return (*msqv1node)(atomic.LoadPointer(p))
    28  }
    29  
    30  func (q *MSQueue) Enqueue(value uint64) bool {
    31  	node := &msqv1node{value: value}
    32  	for {
    33  		tail := atomic.LoadPointer(&q.tail)
    34  		tailstruct := (*msqv1node)(tail)
    35  		next := atomic.LoadPointer(&tailstruct.next)
    36  		if tail == atomic.LoadPointer(&q.tail) {
    37  			if next == nil {
    38  				// tail.next is empty, inset new node
    39  				if atomic.CompareAndSwapPointer(&tailstruct.next, next, unsafe.Pointer(node)) {
    40  					atomic.CompareAndSwapPointer(&q.tail, tail, unsafe.Pointer(node))
    41  					break
    42  				}
    43  			} else {
    44  				atomic.CompareAndSwapPointer(&q.tail, tail, next)
    45  			}
    46  		}
    47  	}
    48  	return true
    49  }
    50  
    51  func (q *MSQueue) Dequeue() (value uint64, ok bool) {
    52  	for {
    53  		head := atomic.LoadPointer(&q.head)
    54  		tail := atomic.LoadPointer(&q.tail)
    55  		headstruct := (*msqv1node)(head)
    56  		next := atomic.LoadPointer(&headstruct.next)
    57  		if head == atomic.LoadPointer(&q.head) {
    58  			if head == tail {
    59  				if next == nil {
    60  					return 0, false
    61  				}
    62  				atomic.CompareAndSwapPointer(&q.tail, tail, next)
    63  			} else {
    64  				value = ((*msqv1node)(next)).value
    65  				if atomic.CompareAndSwapPointer(&q.head, head, next) {
    66  					return value, true
    67  				}
    68  			}
    69  		}
    70  	}
    71  }