github.com/fufuok/utils@v1.0.10/xsync/mpmcqueueof.go (about)

     1  //go:build go1.19
     2  // +build go1.19
     3  
     4  package xsync
     5  
     6  import (
     7  	"runtime"
     8  	"sync/atomic"
     9  	"unsafe"
    10  )
    11  
    12  // A MPMCQueueOf is a bounded multi-producer multi-consumer concurrent
    13  // queue. It's a generic version of MPMCQueue.
    14  //
    15  // MPMCQueue instances must be created with NewMPMCQueueOf function.
    16  // A MPMCQueueOf must not be copied after first use.
    17  //
    18  // Based on the data structure from the following C++ library:
    19  // https://github.com/rigtorp/MPMCQueue
    20  type MPMCQueueOf[I any] struct {
    21  	cap  uint64
    22  	head uint64
    23  	//lint:ignore U1000 prevents false sharing
    24  	hpad [cacheLineSize - 8]byte
    25  	tail uint64
    26  	//lint:ignore U1000 prevents false sharing
    27  	tpad  [cacheLineSize - 8]byte
    28  	slots []slotOfPadded[I]
    29  }
    30  
    31  type slotOfPadded[I any] struct {
    32  	slotOf[I]
    33  	// Unfortunately, proper padding like the below one:
    34  	//
    35  	// pad [cacheLineSize - (unsafe.Sizeof(slotOf[I]{}) % cacheLineSize)]byte
    36  	//
    37  	// won't compile, so here we add a best-effort padding for items up to
    38  	// 56 bytes size.
    39  	//lint:ignore U1000 prevents false sharing
    40  	pad [cacheLineSize - unsafe.Sizeof(atomic.Uint64{})]byte
    41  }
    42  
    43  type slotOf[I any] struct {
    44  	// atomic.Uint64 is used here to get proper 8 byte alignment on
    45  	// 32-bit archs.
    46  	turn atomic.Uint64
    47  	item I
    48  }
    49  
    50  // NewMPMCQueueOf creates a new MPMCQueueOf instance with the given
    51  // capacity.
    52  func NewMPMCQueueOf[I any](capacity int) *MPMCQueueOf[I] {
    53  	if capacity < 1 {
    54  		panic("capacity must be positive number")
    55  	}
    56  	return &MPMCQueueOf[I]{
    57  		cap:   uint64(capacity),
    58  		slots: make([]slotOfPadded[I], capacity),
    59  	}
    60  }
    61  
    62  // Enqueue inserts the given item into the queue.
    63  // Blocks, if the queue is full.
    64  func (q *MPMCQueueOf[I]) Enqueue(item I) {
    65  	head := atomic.AddUint64(&q.head, 1) - 1
    66  	slot := &q.slots[q.idx(head)]
    67  	turn := q.turn(head) * 2
    68  	for slot.turn.Load() != turn {
    69  		runtime.Gosched()
    70  	}
    71  	slot.item = item
    72  	slot.turn.Store(turn + 1)
    73  }
    74  
    75  // Dequeue retrieves and removes the item from the head of the queue.
    76  // Blocks, if the queue is empty.
    77  func (q *MPMCQueueOf[I]) Dequeue() I {
    78  	var zeroedI I
    79  	tail := atomic.AddUint64(&q.tail, 1) - 1
    80  	slot := &q.slots[q.idx(tail)]
    81  	turn := q.turn(tail)*2 + 1
    82  	for slot.turn.Load() != turn {
    83  		runtime.Gosched()
    84  	}
    85  	item := slot.item
    86  	slot.item = zeroedI
    87  	slot.turn.Store(turn + 1)
    88  	return item
    89  }
    90  
    91  // TryEnqueue inserts the given item into the queue. Does not block
    92  // and returns immediately. The result indicates that the queue isn't
    93  // full and the item was inserted.
    94  func (q *MPMCQueueOf[I]) TryEnqueue(item I) bool {
    95  	head := atomic.LoadUint64(&q.head)
    96  	for {
    97  		slot := &q.slots[q.idx(head)]
    98  		turn := q.turn(head) * 2
    99  		if slot.turn.Load() == turn {
   100  			if atomic.CompareAndSwapUint64(&q.head, head, head+1) {
   101  				slot.item = item
   102  				slot.turn.Store(turn + 1)
   103  				return true
   104  			}
   105  		} else {
   106  			prevHead := head
   107  			head = atomic.LoadUint64(&q.head)
   108  			if head == prevHead {
   109  				return false
   110  			}
   111  		}
   112  		runtime.Gosched()
   113  	}
   114  }
   115  
   116  // TryDequeue retrieves and removes the item from the head of the
   117  // queue. Does not block and returns immediately. The ok result
   118  // indicates that the queue isn't empty and an item was retrieved.
   119  func (q *MPMCQueueOf[I]) TryDequeue() (item I, ok bool) {
   120  	tail := atomic.LoadUint64(&q.tail)
   121  	for {
   122  		slot := &q.slots[q.idx(tail)]
   123  		turn := q.turn(tail)*2 + 1
   124  		if slot.turn.Load() == turn {
   125  			if atomic.CompareAndSwapUint64(&q.tail, tail, tail+1) {
   126  				var zeroedI I
   127  				item = slot.item
   128  				ok = true
   129  				slot.item = zeroedI
   130  				slot.turn.Store(turn + 1)
   131  				return
   132  			}
   133  		} else {
   134  			prevTail := tail
   135  			tail = atomic.LoadUint64(&q.tail)
   136  			if tail == prevTail {
   137  				return
   138  			}
   139  		}
   140  		runtime.Gosched()
   141  	}
   142  }
   143  
   144  func (q *MPMCQueueOf[I]) idx(i uint64) uint64 {
   145  	return i % q.cap
   146  }
   147  
   148  func (q *MPMCQueueOf[I]) turn(i uint64) uint64 {
   149  	return i / q.cap
   150  }