github.com/ergo-services/ergo@v1.999.224/lib/mpsc.go (about)

     1  // this is a lock-free implementation of MPSC queue (Multiple Producers Single Consumer)
     2  
     3  package lib
     4  
     5  import (
     6  	"math"
     7  	"sync/atomic"
     8  	"unsafe"
     9  )
    10  
    11  type queueMPSC struct {
    12  	head *itemMPSC
    13  	tail *itemMPSC
    14  }
    15  
    16  type queueLimitMPSC struct {
    17  	head   *itemMPSC
    18  	tail   *itemMPSC
    19  	length int64
    20  	limit  int64
    21  }
    22  
    23  type QueueMPSC interface {
    24  	Push(value interface{}) bool
    25  	Pop() (interface{}, bool)
    26  	Item() ItemMPSC
    27  	// Len returns the number of items in the queue
    28  	Len() int64
    29  }
    30  
    31  func NewQueueMPSC() QueueMPSC {
    32  	emptyItem := &itemMPSC{}
    33  	return &queueMPSC{
    34  		head: emptyItem,
    35  		tail: emptyItem,
    36  	}
    37  }
    38  
    39  func NewQueueLimitMPSC(limit int64) QueueMPSC {
    40  	if limit < 1 {
    41  		limit = math.MaxInt64
    42  	}
    43  	emptyItem := &itemMPSC{}
    44  	return &queueLimitMPSC{
    45  		limit: limit,
    46  		head:  emptyItem,
    47  		tail:  emptyItem,
    48  	}
    49  }
    50  
    51  type ItemMPSC interface {
    52  	Next() ItemMPSC
    53  	Value() interface{}
    54  	Clear()
    55  }
    56  
    57  type itemMPSC struct {
    58  	value interface{}
    59  	next  *itemMPSC
    60  }
    61  
    62  // Push place the given value in the queue head (FIFO). Returns always true
    63  func (q *queueMPSC) Push(value interface{}) bool {
    64  	i := &itemMPSC{
    65  		value: value,
    66  	}
    67  	old_head := (*itemMPSC)(atomic.SwapPointer((*unsafe.Pointer)(unsafe.Pointer(&q.head)), unsafe.Pointer(i)))
    68  	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&old_head.next)), unsafe.Pointer(i))
    69  	return true
    70  }
    71  
    72  // Push place the given value in the queue head (FIFO). Returns false if exceeded the limit
    73  func (q *queueLimitMPSC) Push(value interface{}) bool {
    74  	if q.Len()+1 > q.limit {
    75  		return false
    76  	}
    77  	atomic.AddInt64(&q.length, 1)
    78  	i := &itemMPSC{
    79  		value: value,
    80  	}
    81  	old_head := (*itemMPSC)(atomic.SwapPointer((*unsafe.Pointer)(unsafe.Pointer(&q.head)), unsafe.Pointer(i)))
    82  	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&old_head.next)), unsafe.Pointer(i))
    83  	return true
    84  }
    85  
    86  // Pop takes the item from the queue tail. Returns false if the queue is empty. Can be used in a single consumer (goroutine) only.
    87  func (q *queueMPSC) Pop() (interface{}, bool) {
    88  	tail_next := (*itemMPSC)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&q.tail.next))))
    89  	if tail_next == nil {
    90  		return nil, false
    91  	}
    92  
    93  	value := tail_next.value
    94  	tail_next.value = nil // let the GC free this item
    95  	q.tail = tail_next
    96  	return value, true
    97  }
    98  
    99  // Pop takes the item from the queue tail. Returns false if the queue is empty. Can be used in a single consumer (goroutine) only.
   100  func (q *queueLimitMPSC) Pop() (interface{}, bool) {
   101  	tail_next := (*itemMPSC)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&q.tail.next))))
   102  	if tail_next == nil {
   103  		return nil, false
   104  	}
   105  
   106  	value := tail_next.value
   107  	tail_next.value = nil // let the GC free this item
   108  	q.tail = tail_next
   109  	atomic.AddInt64(&q.length, -1)
   110  	return value, true
   111  }
   112  
   113  // Len returns -1 for the queue with no limit
   114  func (q *queueMPSC) Len() int64 {
   115  	return -1
   116  }
   117  
   118  // Len returns queue length
   119  func (q *queueLimitMPSC) Len() int64 {
   120  	return atomic.LoadInt64(&q.length)
   121  }
   122  
   123  // Item returns the tail item of the queue. Returns nil if queue is empty.
   124  func (q *queueMPSC) Item() ItemMPSC {
   125  	item := (*itemMPSC)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&q.tail.next))))
   126  	if item == nil {
   127  		return nil
   128  	}
   129  	return item
   130  }
   131  
   132  // Item returns the tail item of the queue. Returns nil if queue is empty.
   133  func (q *queueLimitMPSC) Item() ItemMPSC {
   134  	item := (*itemMPSC)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&q.tail.next))))
   135  	if item == nil {
   136  		return nil
   137  	}
   138  	return item
   139  }
   140  
   141  //
   142  // ItemMPSC interface implementation
   143  //
   144  
   145  // Next provides walking through the queue. Returns nil if the last item is reached.
   146  func (i *itemMPSC) Next() ItemMPSC {
   147  	next := (*itemMPSC)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&i.next))))
   148  	if next == nil {
   149  		return nil
   150  	}
   151  	return next
   152  }
   153  
   154  // Value returns stored value of the queue item
   155  func (i *itemMPSC) Value() interface{} {
   156  	return i.value
   157  }
   158  
   159  // Clear sets the value to nil. It doesn't remove this item from the queue. Can be used in a signle consumer (goroutine) only.
   160  func (i *itemMPSC) Clear() {
   161  	i.value = nil
   162  }