github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/queue/ring_buffer.go (about)

     1  package queue
     2  
     3  // circular queue
     4  
     5  import (
     6  	"runtime"
     7  	"sync/atomic"
     8  
     9  	"github.com/benz9527/toy-box/algo/bit"
    10  )
    11  
    12  var (
    13  	_ RingBufferEntry[struct{}] = (*rbEntry[struct{}])(nil)
    14  	_ RingBuffer[struct{}]      = (*xRingBuffer[struct{}])(nil)
    15  )
    16  
    17  type rbEntry[T any] struct {
    18  	// index of the element in the ring buffer and
    19  	// it is also the "lock" to protect the value
    20  	// in lock-free mode by atomic operation
    21  	cursor uint64
    22  	value  T
    23  }
    24  
    25  func (e *rbEntry[T]) GetValue() T {
    26  	return e.value
    27  }
    28  
    29  func (e *rbEntry[T]) GetCursor() uint64 {
    30  	return atomic.LoadUint64(&e.cursor)
    31  }
    32  
    33  func (e *rbEntry[T]) Store(cursor uint64, value T) {
    34  	// atomic operation should be called at the end of the function
    35  	// otherwise, the value of cursor may be changed by other goroutines
    36  	// Go1.19 atomic guarantees the sequentially consistent order
    37  	// Go1.18 atomic.Pointer[T]
    38  	e.value = value
    39  	atomic.StoreUint64(&e.cursor, cursor)
    40  }
    41  
    42  type xRingBuffer[T any] struct {
    43  	capacityMask uint64
    44  	buffer       []RingBufferEntry[T]
    45  }
    46  
    47  const _100M = 100 * 1024 * 1024
    48  
    49  func NewXRingBuffer[T any](capacity uint64) RingBuffer[T] {
    50  	if capacity > _100M {
    51  		panic("capacity is too large")
    52  	}
    53  	if bit.IsPowOf2(capacity) {
    54  		capacity = bit.RoundupPowOf2ByCeil(capacity)
    55  		if capacity > _100M {
    56  			panic("capacity is too large")
    57  		}
    58  	}
    59  
    60  	rb := &xRingBuffer[T]{
    61  		capacityMask: capacity - 1,
    62  		buffer:       make([]RingBufferEntry[T], capacity),
    63  	}
    64  	for i := uint64(0); i < capacity; i++ {
    65  		rb.buffer[i] = &rbEntry[T]{}
    66  	}
    67  	runtime.SetFinalizer(rb, func(rb *xRingBuffer[T]) {
    68  		clear(rb.buffer)
    69  	})
    70  	return rb
    71  }
    72  
    73  func (rb *xRingBuffer[T]) Capacity() uint64 {
    74  	return rb.capacityMask + 1
    75  }
    76  
    77  func (rb *xRingBuffer[T]) LoadEntryByCursor(cursor uint64) RingBufferEntry[T] {
    78  	return rb.buffer[cursor&rb.capacityMask]
    79  }