github.com/rudderlabs/rudder-go-kit@v0.30.0/queue/priorityQueue.go (about)

     1  package queue
     2  
     3  import (
     4  	"container/heap"
     5  	"time"
     6  )
     7  
     8  // Item stores the attributes which will be pushed to the priority queue..
     9  type Item[T any] struct {
    10  	Value     T
    11  	Priority  int
    12  	timeStamp int64
    13  	index     int
    14  }
    15  
    16  // PriorityQueue provides a heap.Interface compatible priority queue for the Item type.
    17  // The actual Item.index in the queue is controlled by the Item.Priority and Item.timeStamp.
    18  type PriorityQueue[T any] []*Item[T]
    19  
    20  // Len: Size of the priority queue . Used to satisfy the heap interface...
    21  func (pq PriorityQueue[T]) Len() int { return len(pq) }
    22  
    23  // Less is used to compare elements and store them in the proper order in
    24  // priority queue.
    25  func (pq PriorityQueue[T]) Less(i, j int) bool {
    26  	if pq[i].Priority == pq[j].Priority {
    27  		return pq[i].timeStamp <= pq[j].timeStamp
    28  	}
    29  	return pq[i].Priority > pq[j].Priority
    30  }
    31  
    32  // Swap is used to swap the values in the priority queue.
    33  func (pq PriorityQueue[T]) Swap(i, j int) {
    34  	pq[i], pq[j] = pq[j], pq[i]
    35  	pq[i].index = i
    36  	pq[j].index = j
    37  }
    38  
    39  // Push adds elements to the priority queue
    40  func (pq *PriorityQueue[T]) Push(x interface{}) {
    41  	n := len(*pq)
    42  	item := x.(*Item[T])
    43  	item.index = n
    44  	item.timeStamp = makeTimestamp()
    45  	*pq = append(*pq, item)
    46  }
    47  
    48  // Pop removes elements from the priority queue
    49  func (pq *PriorityQueue[T]) Pop() interface{} {
    50  	old := *pq
    51  	n := len(old)
    52  	if n == 0 {
    53  		return nil
    54  	}
    55  	item := old[n-1]
    56  	old[n-1] = nil
    57  	item.index = -1
    58  	*pq = old[0 : n-1]
    59  	return item
    60  }
    61  
    62  // Top returns the topmost element in the priority queue
    63  func (pq *PriorityQueue[T]) Top() interface{} {
    64  	if len(*pq) == 0 {
    65  		return nil
    66  	}
    67  	ol := *pq
    68  	return *ol[0]
    69  }
    70  
    71  // GetIndex returns the index of the corresponding element.
    72  func (pq *PriorityQueue[T]) GetIndex(x interface{}) int {
    73  	item := x.(*Item[T])
    74  	return item.index
    75  }
    76  
    77  // Update updates the attributes of an element in the priority queue.
    78  func (pq *PriorityQueue[T]) Update(item *Item[T], priority int) {
    79  	if item.index == -1 {
    80  		return
    81  	}
    82  	item.Priority = priority
    83  	heap.Fix(pq, item.index)
    84  }
    85  
    86  func makeTimestamp() int64 {
    87  	return time.Now().UnixNano() / int64(time.Millisecond)
    88  }