github.com/hashicorp/vault/sdk@v0.13.0/queue/priority_queue.go (about) 1 // Copyright (c) HashiCorp, Inc. 2 // SPDX-License-Identifier: MPL-2.0 3 4 // Package queue provides Vault plugins with a Priority Queue. It can be used 5 // as an in-memory list of queue.Item sorted by their priority, and offers 6 // methods to find or remove items by their key. Internally it uses 7 // container/heap; see Example Priority Queue: 8 // https://golang.org/pkg/container/heap/#example__priorityQueue 9 package queue 10 11 import ( 12 "container/heap" 13 "errors" 14 "sync" 15 16 "github.com/mitchellh/copystructure" 17 ) 18 19 // ErrEmpty is returned for queues with no items 20 var ErrEmpty = errors.New("queue is empty") 21 22 // ErrDuplicateItem is returned when the queue attmepts to push an item to a key that 23 // already exists. The queue does not attempt to update, instead returns this 24 // error. If an Item needs to be updated or replaced, pop the item first. 25 var ErrDuplicateItem = errors.New("duplicate item") 26 27 // New initializes the internal data structures and returns a new 28 // PriorityQueue 29 func New() *PriorityQueue { 30 pq := PriorityQueue{ 31 data: make(queue, 0), 32 dataMap: make(map[string]*Item), 33 } 34 heap.Init(&pq.data) 35 return &pq 36 } 37 38 // PriorityQueue facilitates queue of Items, providing Push, Pop, and 39 // PopByKey convenience methods. The ordering (priority) is an int64 value 40 // with the smallest value is the highest priority. PriorityQueue maintains both 41 // an internal slice for the queue as well as a map of the same items with their 42 // keys as the index. This enables users to find specific items by key. The map 43 // must be kept in sync with the data slice. 44 // See https://golang.org/pkg/container/heap/#example__priorityQueue 45 type PriorityQueue struct { 46 // data is the internal structure that holds the queue, and is operated on by 47 // heap functions 48 data queue 49 50 // dataMap represents all the items in the queue, with unique indexes, used 51 // for finding specific items. dataMap is kept in sync with the data slice 52 dataMap map[string]*Item 53 54 // lock is a read/write mutex, and used to facilitate read/write locks on the 55 // data and dataMap fields 56 lock sync.RWMutex 57 } 58 59 // queue is the internal data structure used to satisfy heap.Interface. This 60 // prevents users from calling Pop and Push heap methods directly 61 type queue []*Item 62 63 // Item is something managed in the priority queue 64 type Item struct { 65 // Key is a unique string used to identify items in the internal data map 66 Key string 67 // Value is an unspecified type that implementations can use to store 68 // information 69 Value interface{} 70 71 // Priority determines ordering in the queue, with the lowest value being the 72 // highest priority 73 Priority int64 74 75 // index is an internal value used by the heap package, and should not be 76 // modified by any consumer of the priority queue 77 index int 78 } 79 80 // Len returns the count of items in the Priority Queue 81 func (pq *PriorityQueue) Len() int { 82 pq.lock.RLock() 83 defer pq.lock.RUnlock() 84 return pq.data.Len() 85 } 86 87 // Pop pops the highest priority item from the queue. This is a 88 // wrapper/convenience method that calls heap.Pop, so consumers do not need to 89 // invoke heap functions directly 90 func (pq *PriorityQueue) Pop() (*Item, error) { 91 pq.lock.Lock() 92 defer pq.lock.Unlock() 93 94 if pq.data.Len() == 0 { 95 return nil, ErrEmpty 96 } 97 98 item := heap.Pop(&pq.data).(*Item) 99 delete(pq.dataMap, item.Key) 100 return item, nil 101 } 102 103 // Push pushes an item on to the queue. This is a wrapper/convenience 104 // method that calls heap.Push, so consumers do not need to invoke heap 105 // functions directly. Items must have unique Keys, and Items in the queue 106 // cannot be updated. To modify an Item, users must first remove it and re-push 107 // it after modifications 108 func (pq *PriorityQueue) Push(i *Item) error { 109 if i == nil || i.Key == "" { 110 return errors.New("error adding item: Item Key is required") 111 } 112 113 pq.lock.Lock() 114 defer pq.lock.Unlock() 115 116 if _, ok := pq.dataMap[i.Key]; ok { 117 return ErrDuplicateItem 118 } 119 // Copy the item value(s) so that modifications to the source item does not 120 // affect the item on the queue 121 clone, err := copystructure.Copy(i) 122 if err != nil { 123 return err 124 } 125 126 pq.dataMap[i.Key] = clone.(*Item) 127 heap.Push(&pq.data, clone) 128 return nil 129 } 130 131 // PopByKey searches the queue for an item with the given key and removes it 132 // from the queue if found. Returns nil if not found. This method must fix the 133 // queue after removing any key. 134 func (pq *PriorityQueue) PopByKey(key string) (*Item, error) { 135 pq.lock.Lock() 136 defer pq.lock.Unlock() 137 138 item, ok := pq.dataMap[key] 139 if !ok { 140 return nil, nil 141 } 142 143 // Remove the item the heap and delete it from the dataMap 144 itemRaw := heap.Remove(&pq.data, item.index) 145 delete(pq.dataMap, key) 146 147 if itemRaw != nil { 148 if i, ok := itemRaw.(*Item); ok { 149 return i, nil 150 } 151 } 152 153 return nil, nil 154 } 155 156 // Len returns the number of items in the queue data structure. Do not use this 157 // method directly on the queue, use PriorityQueue.Len() instead. 158 func (q queue) Len() int { return len(q) } 159 160 // Less returns whether the Item with index i should sort before the Item with 161 // index j in the queue. This method is used by the queue to determine priority 162 // internally; the Item with the lower value wins. (priority zero is higher 163 // priority than 1). The priority of Items with equal values is undetermined. 164 func (q queue) Less(i, j int) bool { 165 return q[i].Priority < q[j].Priority 166 } 167 168 // Swap swaps things in-place; part of sort.Interface 169 func (q queue) Swap(i, j int) { 170 q[i], q[j] = q[j], q[i] 171 q[i].index = i 172 q[j].index = j 173 } 174 175 // Push is used by heap.Interface to push items onto the heap. This method is 176 // invoked by container/heap, and should not be used directly. 177 // See: https://golang.org/pkg/container/heap/#Interface 178 func (q *queue) Push(x interface{}) { 179 n := len(*q) 180 item := x.(*Item) 181 item.index = n 182 *q = append(*q, item) 183 } 184 185 // Pop is used by heap.Interface to pop items off of the heap. This method is 186 // invoked by container/heap, and should not be used directly. 187 // See: https://golang.org/pkg/container/heap/#Interface 188 func (q *queue) Pop() interface{} { 189 old := *q 190 n := len(old) 191 item := old[n-1] 192 old[n-1] = nil // avoid memory leak 193 item.index = -1 // for safety 194 *q = old[0 : n-1] 195 return item 196 }