github.com/moontrade/nogc@v0.1.7/collections/queue/lock_free.go (about)

     1  //go:build !tinygo
     2  // +build !tinygo
     3  
     4  package queue
     5  
     6  import (
     7  	"github.com/moontrade/nogc"
     8  	"sync/atomic"
     9  	"unsafe"
    10  )
    11  
    12  // LockFree is a simple non-blocking and concurrent queue.
    13  type LockFree struct {
    14  	head   uintptr
    15  	tail   uintptr
    16  	length int32
    17  }
    18  
    19  type lockFreeNode struct {
    20  	value nogc.Bytes
    21  	next  uintptr
    22  }
    23  
    24  //goland:noinspection GoVetUnsafePointer
    25  func AllocLockFreeQueue() *LockFree {
    26  	q := (*LockFree)(unsafe.Pointer(nogc.Alloc(unsafe.Sizeof(LockFree{}))))
    27  	n := nogc.Alloc(unsafe.Sizeof(lockFreeNode{}))
    28  	node := (*lockFreeNode)(unsafe.Pointer(n))
    29  	node.value = nogc.Bytes{}
    30  	node.next = 0
    31  	q.head = uintptr(n)
    32  	q.tail = uintptr(n)
    33  	q.length = 0
    34  	return q
    35  }
    36  
    37  func (l *LockFree) Free() {
    38  	nogc.Free(nogc.Pointer(unsafe.Pointer(l)))
    39  }
    40  
    41  // Enqueue puts the given value v at the tail of the queue.
    42  //goland:noinspection GoVetUnsafePointer
    43  func (q *LockFree) Enqueue(task nogc.Bytes) {
    44  	n := uintptr(nogc.Alloc(unsafe.Sizeof(lockFreeNode{})))
    45  	node := (*lockFreeNode)(unsafe.Pointer(n))
    46  	node.value = task
    47  	node.next = 0
    48  retry:
    49  	last := atomic.LoadUintptr(&q.tail)
    50  	lastV := (*lockFreeNode)(unsafe.Pointer(last))
    51  	next := atomic.LoadUintptr(&lastV.next)
    52  	// Are tail and next consistent?
    53  	if last == atomic.LoadUintptr(&q.tail) {
    54  		if next == 0 {
    55  			// Try to link node at the end of the linked list.
    56  			if atomic.CompareAndSwapUintptr(&lastV.next, next, n) { // enqueue is done.
    57  				// Try to swing tail to the inserted node.
    58  				atomic.CompareAndSwapUintptr(&q.tail, last, n)
    59  				atomic.AddInt32(&q.length, 1)
    60  				return
    61  			}
    62  		} else { // tail was not pointing to the last node
    63  			// Try to swing tail to the next node.
    64  			atomic.CompareAndSwapUintptr(&q.tail, last, next)
    65  		}
    66  	}
    67  	goto retry
    68  }
    69  
    70  // Dequeue removes and returns the value at the head of the queue.
    71  // It returns nil if the queue is empty.
    72  //goland:noinspection GoVetUnsafePointer
    73  func (q *LockFree) Dequeue() nogc.Bytes {
    74  retry:
    75  	first := atomic.LoadUintptr(&q.head)
    76  	firstV := (*lockFreeNode)(unsafe.Pointer(first))
    77  	last := atomic.LoadUintptr(&q.tail)
    78  	next := atomic.LoadUintptr(&firstV.next)
    79  	// Are first, tail, and next consistent?
    80  	if first == atomic.LoadUintptr(&q.head) {
    81  		// Is queue empty or tail falling behind?
    82  		if first == last {
    83  			// Is queue empty?
    84  			if next == 0 {
    85  				//println("empty")
    86  				return nogc.Bytes{}
    87  			}
    88  			//println("first == tail")
    89  			atomic.CompareAndSwapUintptr(&q.tail, last, next) // tail is falling behind, try to advance it.
    90  		} else {
    91  			// Read value before CAS, otherwise another dequeue might free the next node.
    92  			task := (*lockFreeNode)(unsafe.Pointer(next)).value
    93  			if atomic.CompareAndSwapUintptr(&q.head, first, next) { // dequeue is done, return value.
    94  				atomic.AddInt32(&q.length, -1)
    95  				nogc.Free(nogc.Pointer(first))
    96  				return task
    97  			}
    98  		}
    99  	}
   100  	goto retry
   101  }
   102  
   103  // Empty indicates whether this queue is empty or not.
   104  func (q *LockFree) Empty() bool {
   105  	return atomic.LoadInt32(&q.length) == 0
   106  }