gitee.com/lh-her-team/common@v1.5.1/queue/lockfreequeue/annular_lockfree_queue.go (about)

     1  package lockfreequeue
     2  
     3  import (
     4  	"fmt"
     5  	"runtime"
     6  	"sync/atomic"
     7  )
     8  
     9  type valuePit struct {
    10  	val  interface{}
    11  	rIdx uint32
    12  	wIdx uint32
    13  }
    14  
    15  type Queue struct {
    16  	capacity  uint32
    17  	capMod    uint32
    18  	readerIdx uint32
    19  	writerIdx uint32
    20  	pits      []*valuePit
    21  }
    22  
    23  func NewQueue(cap uint32) *Queue {
    24  	q := new(Queue)
    25  	q.capacity = minQuantity(cap)
    26  	q.capMod = q.capacity - 1
    27  	q.readerIdx = 0
    28  	q.writerIdx = 0
    29  	q.pits = make([]*valuePit, q.capacity)
    30  	for i := range q.pits {
    31  		q.pits[i] = new(valuePit)
    32  	}
    33  	return q
    34  }
    35  
    36  func (q *Queue) String() string {
    37  	rIdx := atomic.LoadUint32(&q.readerIdx)
    38  	wIdx := atomic.LoadUint32(&q.writerIdx)
    39  	return fmt.Sprintf("Queue{capacity: %d, capMod: %d, readerIdx: %d, writerIdx: %d}", q.capacity, q.capMod, rIdx, wIdx)
    40  }
    41  
    42  func (q *Queue) Capacity() uint32 {
    43  	return q.capacity
    44  }
    45  
    46  func (q *Queue) Quantity() uint32 {
    47  	var wIdx, rIdx, currentQuantity uint32
    48  	wIdx = atomic.LoadUint32(&q.writerIdx)
    49  	rIdx = atomic.LoadUint32(&q.readerIdx)
    50  	if rIdx > wIdx {
    51  		currentQuantity = 0
    52  	} else {
    53  		currentQuantity = wIdx - rIdx
    54  	}
    55  	return currentQuantity
    56  }
    57  
    58  func (q *Queue) Push(val interface{}) (ok bool, quantity uint32) {
    59  	var wIdxNew, currentQuantity uint32
    60  	capMod := q.capMod
    61  	currentQuantity = q.Quantity()
    62  	if currentQuantity >= q.capacity {
    63  		return false, currentQuantity
    64  	}
    65  	wIdxNew = atomic.AddUint32(&q.writerIdx, 1)
    66  	idx := (wIdxNew - 1) & capMod
    67  	pit := q.pits[idx]
    68  	for {
    69  		pwIdx := atomic.LoadUint32(&pit.wIdx)
    70  		prIdx := atomic.LoadUint32(&pit.rIdx)
    71  		if pwIdx == prIdx && (wIdxNew == pwIdx+q.capacity || pwIdx == 0) {
    72  			pit.val = val
    73  			atomic.StoreUint32(&pit.wIdx, wIdxNew)
    74  			return true, currentQuantity + 1
    75  		}
    76  		runtime.Gosched()
    77  	}
    78  }
    79  
    80  func (q *Queue) Pull() (val interface{}, ok bool, quantity uint32) {
    81  	var rIdxNew, currentQuantity uint32
    82  	capMod := q.capMod
    83  	currentQuantity = q.Quantity()
    84  	if currentQuantity < 1 {
    85  		return nil, false, currentQuantity
    86  	}
    87  	rIdxNew = atomic.AddUint32(&q.readerIdx, 1)
    88  	idx := (rIdxNew - 1) & capMod
    89  	pit := q.pits[idx]
    90  	for {
    91  		pwIdx := atomic.LoadUint32(&pit.wIdx)
    92  		prIdx := atomic.LoadUint32(&pit.rIdx)
    93  		if rIdxNew == pwIdx && (rIdxNew == prIdx+q.capacity || prIdx == 0) {
    94  			val = pit.val
    95  			pit.val = nil
    96  			atomic.StoreUint32(&pit.rIdx, rIdxNew)
    97  			return val, true, currentQuantity - 1
    98  		}
    99  		runtime.Gosched()
   100  	}
   101  }
   102  
   103  // round 到最近的2的幂值
   104  func minQuantity(v uint32) uint32 {
   105  	v--
   106  	v |= v >> 1
   107  	v |= v >> 2
   108  	v |= v >> 4
   109  	v |= v >> 8
   110  	v |= v >> 16
   111  	v++
   112  	return v
   113  }