github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/priority_queue.go (about)

     1  package util
     2  
     3  import (
     4  	"container/heap"
     5  	"sync"
     6  
     7  	"github.com/prometheus/client_golang/prometheus"
     8  )
     9  
    10  // PriorityQueue is a priority queue.
    11  type PriorityQueue struct {
    12  	lock        sync.Mutex
    13  	cond        *sync.Cond
    14  	closing     bool
    15  	closed      bool
    16  	hit         map[string]struct{}
    17  	queue       queue
    18  	lengthGauge prometheus.Gauge
    19  }
    20  
    21  // Op is an operation on the priority queue.
    22  type Op interface {
    23  	Key() string
    24  	Priority() int64 // The larger the number the higher the priority.
    25  }
    26  
    27  type queue []Op
    28  
    29  func (q queue) Len() int           { return len(q) }
    30  func (q queue) Less(i, j int) bool { return q[i].Priority() > q[j].Priority() }
    31  func (q queue) Swap(i, j int)      { q[i], q[j] = q[j], q[i] }
    32  
    33  // Push and Pop use pointer receivers because they modify the slice's length,
    34  // not just its contents.
    35  func (q *queue) Push(x interface{}) {
    36  	*q = append(*q, x.(Op))
    37  }
    38  
    39  func (q *queue) Pop() interface{} {
    40  	old := *q
    41  	n := len(old)
    42  	x := old[n-1]
    43  	*q = old[0 : n-1]
    44  	return x
    45  }
    46  
    47  // NewPriorityQueue makes a new priority queue.
    48  func NewPriorityQueue(lengthGauge prometheus.Gauge) *PriorityQueue {
    49  	pq := &PriorityQueue{
    50  		hit:         map[string]struct{}{},
    51  		lengthGauge: lengthGauge,
    52  	}
    53  	pq.cond = sync.NewCond(&pq.lock)
    54  	heap.Init(&pq.queue)
    55  	return pq
    56  }
    57  
    58  // Length returns the length of the queue.
    59  func (pq *PriorityQueue) Length() int {
    60  	pq.lock.Lock()
    61  	defer pq.lock.Unlock()
    62  	return len(pq.queue)
    63  }
    64  
    65  // Close signals that the queue should be closed when it is empty.
    66  // A closed queue will not accept new items.
    67  func (pq *PriorityQueue) Close() {
    68  	pq.lock.Lock()
    69  	defer pq.lock.Unlock()
    70  	pq.closing = true
    71  	pq.cond.Broadcast()
    72  }
    73  
    74  // DiscardAndClose closes the queue and removes all the items from it.
    75  func (pq *PriorityQueue) DiscardAndClose() {
    76  	pq.lock.Lock()
    77  	defer pq.lock.Unlock()
    78  	pq.closed = true
    79  	pq.queue = nil
    80  	pq.hit = map[string]struct{}{}
    81  	pq.cond.Broadcast()
    82  }
    83  
    84  // Enqueue adds an operation to the queue in priority order. Returns
    85  // true if added; false if the operation was already on the queue.
    86  func (pq *PriorityQueue) Enqueue(op Op) bool {
    87  	pq.lock.Lock()
    88  	defer pq.lock.Unlock()
    89  
    90  	if pq.closed {
    91  		panic("enqueue on closed queue")
    92  	}
    93  
    94  	_, enqueued := pq.hit[op.Key()]
    95  	if enqueued {
    96  		return false
    97  	}
    98  
    99  	pq.hit[op.Key()] = struct{}{}
   100  	heap.Push(&pq.queue, op)
   101  	pq.cond.Broadcast()
   102  	if pq.lengthGauge != nil {
   103  		pq.lengthGauge.Inc()
   104  	}
   105  	return true
   106  }
   107  
   108  // Dequeue will return the op with the highest priority; block if queue is
   109  // empty; returns nil if queue is closed.
   110  func (pq *PriorityQueue) Dequeue() Op {
   111  	pq.lock.Lock()
   112  	defer pq.lock.Unlock()
   113  
   114  	for len(pq.queue) == 0 && !(pq.closing || pq.closed) {
   115  		pq.cond.Wait()
   116  	}
   117  
   118  	if len(pq.queue) == 0 && (pq.closing || pq.closed) {
   119  		pq.closed = true
   120  		return nil
   121  	}
   122  
   123  	op := heap.Pop(&pq.queue).(Op)
   124  	delete(pq.hit, op.Key())
   125  	if pq.lengthGauge != nil {
   126  		pq.lengthGauge.Dec()
   127  	}
   128  	return op
   129  }