github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/ring/buffer.go (about)

     1  package ring
     2  
     3  import (
     4  	"runtime"
     5  	"sync/atomic"
     6  )
     7  
     8  // Implementation is based on
     9  //   http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue
    10  
    11  const (
    12  	Size     = 128 << 10 // ~128000
    13  	Mask     = Size - 1
    14  	CellSize = 4 << 10 // 4kb
    15  )
    16  
    17  type Cell struct {
    18  	sequence int64
    19  	Size     int64
    20  	Data     [CellSize]byte
    21  }
    22  
    23  type Ring struct {
    24  	writeat  int64
    25  	readfrom int64
    26  	buffer   [Size]Cell
    27  }
    28  
    29  func New() *Ring {
    30  	ring := &Ring{}
    31  	for i := range ring.buffer {
    32  		ring.buffer[i].sequence = int64(i)
    33  	}
    34  	return ring
    35  }
    36  
    37  func (ring *Ring) Allocate() *Cell {
    38  	busywait := 16
    39  	for {
    40  		pos := atomic.LoadInt64(&ring.writeat)
    41  		cell := &ring.buffer[pos&Mask]
    42  		seq := atomic.LoadInt64(&cell.sequence)
    43  		if seq-pos == 0 {
    44  			if atomic.CompareAndSwapInt64(&ring.writeat, pos, pos+1) {
    45  				return cell
    46  			}
    47  		}
    48  
    49  		if busywait < 0 {
    50  			runtime.Gosched()
    51  		} else {
    52  			busywait--
    53  		}
    54  	}
    55  }
    56  
    57  func (cell *Cell) Enqueue() { atomic.AddInt64(&cell.sequence, 1) }
    58  
    59  func (ring *Ring) Dequeue() (cell *Cell) {
    60  	busywait := 16
    61  	for {
    62  		pos := atomic.LoadInt64(&ring.readfrom)
    63  		cell = &ring.buffer[pos&Mask]
    64  		seq := atomic.LoadInt64(&cell.sequence)
    65  		dif := seq - (pos + 1)
    66  		if dif == 0 {
    67  			if atomic.CompareAndSwapInt64(&ring.readfrom, pos, pos+1) {
    68  				return cell
    69  			}
    70  		}
    71  
    72  		if busywait < 0 {
    73  			runtime.Gosched()
    74  		} else {
    75  			busywait--
    76  		}
    77  	}
    78  }
    79  
    80  func (cell *Cell) Release() { atomic.AddInt64(&cell.sequence, Size-1) }