github.com/nsqio/nsq@v1.3.0/internal/pqueue/pqueue.go (about)

     1  package pqueue
     2  
     3  import (
     4  	"container/heap"
     5  )
     6  
     7  type Item struct {
     8  	Value    interface{}
     9  	Priority int64
    10  	Index    int
    11  }
    12  
    13  // this is a priority queue as implemented by a min heap
    14  // ie. the 0th element is the *lowest* value
    15  type PriorityQueue []*Item
    16  
    17  func New(capacity int) PriorityQueue {
    18  	return make(PriorityQueue, 0, capacity)
    19  }
    20  
    21  func (pq PriorityQueue) Len() int {
    22  	return len(pq)
    23  }
    24  
    25  func (pq PriorityQueue) Less(i, j int) bool {
    26  	return pq[i].Priority < pq[j].Priority
    27  }
    28  
    29  func (pq PriorityQueue) Swap(i, j int) {
    30  	pq[i], pq[j] = pq[j], pq[i]
    31  	pq[i].Index = i
    32  	pq[j].Index = j
    33  }
    34  
    35  func (pq *PriorityQueue) Push(x interface{}) {
    36  	n := len(*pq)
    37  	c := cap(*pq)
    38  	if n+1 > c {
    39  		npq := make(PriorityQueue, n, c*2)
    40  		copy(npq, *pq)
    41  		*pq = npq
    42  	}
    43  	*pq = (*pq)[0 : n+1]
    44  	item := x.(*Item)
    45  	item.Index = n
    46  	(*pq)[n] = item
    47  }
    48  
    49  func (pq *PriorityQueue) Pop() interface{} {
    50  	n := len(*pq)
    51  	c := cap(*pq)
    52  	if n < (c/2) && c > 25 {
    53  		npq := make(PriorityQueue, n, c/2)
    54  		copy(npq, *pq)
    55  		*pq = npq
    56  	}
    57  	item := (*pq)[n-1]
    58  	item.Index = -1
    59  	*pq = (*pq)[0 : n-1]
    60  	return item
    61  }
    62  
    63  func (pq *PriorityQueue) PeekAndShift(max int64) (*Item, int64) {
    64  	if pq.Len() == 0 {
    65  		return nil, 0
    66  	}
    67  
    68  	item := (*pq)[0]
    69  	if item.Priority > max {
    70  		return nil, item.Priority - max
    71  	}
    72  	heap.Remove(pq, 0)
    73  
    74  	return item, 0
    75  }