cmd/bpa-operator/vendor/k8s.io/client-go/tools/cache/heap.go

   1 /*
   2 Copyright 2017 The Kubernetes Authors.
   3
   4 Licensed under the Apache License, Version 2.0 (the "License");
   5 you may not use this file except in compliance with the License.
   6 You may obtain a copy of the License at
   7
   8     http://www.apache.org/licenses/LICENSE-2.0
   9
  10 Unless required by applicable law or agreed to in writing, software
  11 distributed under the License is distributed on an "AS IS" BASIS,
  12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 See the License for the specific language governing permissions and
  14 limitations under the License.
  15 */
  16
  17 // This file implements a heap data structure.
  18
  19 package cache
  20
  21 import (
  22         "container/heap"
  23         "fmt"
  24         "sync"
  25 )
  26
  27 const (
  28         closedMsg = "heap is closed"
  29 )
  30
  31 type LessFunc func(interface{}, interface{}) bool
  32 type heapItem struct {
  33         obj   interface{} // The object which is stored in the heap.
  34         index int         // The index of the object's key in the Heap.queue.
  35 }
  36
  37 type itemKeyValue struct {
  38         key string
  39         obj interface{}
  40 }
  41
  42 // heapData is an internal struct that implements the standard heap interface
  43 // and keeps the data stored in the heap.
  44 type heapData struct {
  45         // items is a map from key of the objects to the objects and their index.
  46         // We depend on the property that items in the map are in the queue and vice versa.
  47         items map[string]*heapItem
  48         // queue implements a heap data structure and keeps the order of elements
  49         // according to the heap invariant. The queue keeps the keys of objects stored
  50         // in "items".
  51         queue []string
  52
  53         // keyFunc is used to make the key used for queued item insertion and retrieval, and
  54         // should be deterministic.
  55         keyFunc KeyFunc
  56         // lessFunc is used to compare two objects in the heap.
  57         lessFunc LessFunc
  58 }
  59
  60 var (
  61         _ = heap.Interface(&heapData{}) // heapData is a standard heap
  62 )
  63
  64 // Less compares two objects and returns true if the first one should go
  65 // in front of the second one in the heap.
  66 func (h *heapData) Less(i, j int) bool {
  67         if i > len(h.queue) || j > len(h.queue) {
  68                 return false
  69         }
  70         itemi, ok := h.items[h.queue[i]]
  71         if !ok {
  72                 return false
  73         }
  74         itemj, ok := h.items[h.queue[j]]
  75         if !ok {
  76                 return false
  77         }
  78         return h.lessFunc(itemi.obj, itemj.obj)
  79 }
  80
  81 // Len returns the number of items in the Heap.
  82 func (h *heapData) Len() int { return len(h.queue) }
  83
  84 // Swap implements swapping of two elements in the heap. This is a part of standard
  85 // heap interface and should never be called directly.
  86 func (h *heapData) Swap(i, j int) {
  87         h.queue[i], h.queue[j] = h.queue[j], h.queue[i]
  88         item := h.items[h.queue[i]]
  89         item.index = i
  90         item = h.items[h.queue[j]]
  91         item.index = j
  92 }
  93
  94 // Push is supposed to be called by heap.Push only.
  95 func (h *heapData) Push(kv interface{}) {
  96         keyValue := kv.(*itemKeyValue)
  97         n := len(h.queue)
  98         h.items[keyValue.key] = &heapItem{keyValue.obj, n}
  99         h.queue = append(h.queue, keyValue.key)
 100 }
 101
 102 // Pop is supposed to be called by heap.Pop only.
 103 func (h *heapData) Pop() interface{} {
 104         key := h.queue[len(h.queue)-1]
 105         h.queue = h.queue[0 : len(h.queue)-1]
 106         item, ok := h.items[key]
 107         if !ok {
 108                 // This is an error
 109                 return nil
 110         }
 111         delete(h.items, key)
 112         return item.obj
 113 }
 114
 115 // Heap is a thread-safe producer/consumer queue that implements a heap data structure.
 116 // It can be used to implement priority queues and similar data structures.
 117 type Heap struct {
 118         lock sync.RWMutex
 119         cond sync.Cond
 120
 121         // data stores objects and has a queue that keeps their ordering according
 122         // to the heap invariant.
 123         data *heapData
 124
 125         // closed indicates that the queue is closed.
 126         // It is mainly used to let Pop() exit its control loop while waiting for an item.
 127         closed bool
 128 }
 129
 130 // Close the Heap and signals condition variables that may be waiting to pop
 131 // items from the heap.
 132 func (h *Heap) Close() {
 133         h.lock.Lock()
 134         defer h.lock.Unlock()
 135         h.closed = true
 136         h.cond.Broadcast()
 137 }
 138
 139 // Add inserts an item, and puts it in the queue. The item is updated if it
 140 // already exists.
 141 func (h *Heap) Add(obj interface{}) error {
 142         key, err := h.data.keyFunc(obj)
 143         if err != nil {
 144                 return KeyError{obj, err}
 145         }
 146         h.lock.Lock()
 147         defer h.lock.Unlock()
 148         if h.closed {
 149                 return fmt.Errorf(closedMsg)
 150         }
 151         if _, exists := h.data.items[key]; exists {
 152                 h.data.items[key].obj = obj
 153                 heap.Fix(h.data, h.data.items[key].index)
 154         } else {
 155                 h.addIfNotPresentLocked(key, obj)
 156         }
 157         h.cond.Broadcast()
 158         return nil
 159 }
 160
 161 // Adds all the items in the list to the queue and then signals the condition
 162 // variable. It is useful when the caller would like to add all of the items
 163 // to the queue before consumer starts processing them.
 164 func (h *Heap) BulkAdd(list []interface{}) error {
 165         h.lock.Lock()
 166         defer h.lock.Unlock()
 167         if h.closed {
 168                 return fmt.Errorf(closedMsg)
 169         }
 170         for _, obj := range list {
 171                 key, err := h.data.keyFunc(obj)
 172                 if err != nil {
 173                         return KeyError{obj, err}
 174                 }
 175                 if _, exists := h.data.items[key]; exists {
 176                         h.data.items[key].obj = obj
 177                         heap.Fix(h.data, h.data.items[key].index)
 178                 } else {
 179                         h.addIfNotPresentLocked(key, obj)
 180                 }
 181         }
 182         h.cond.Broadcast()
 183         return nil
 184 }
 185
 186 // AddIfNotPresent inserts an item, and puts it in the queue. If an item with
 187 // the key is present in the map, no changes is made to the item.
 188 //
 189 // This is useful in a single producer/consumer scenario so that the consumer can
 190 // safely retry items without contending with the producer and potentially enqueueing
 191 // stale items.
 192 func (h *Heap) AddIfNotPresent(obj interface{}) error {
 193         id, err := h.data.keyFunc(obj)
 194         if err != nil {
 195                 return KeyError{obj, err}
 196         }
 197         h.lock.Lock()
 198         defer h.lock.Unlock()
 199         if h.closed {
 200                 return fmt.Errorf(closedMsg)
 201         }
 202         h.addIfNotPresentLocked(id, obj)
 203         h.cond.Broadcast()
 204         return nil
 205 }
 206
 207 // addIfNotPresentLocked assumes the lock is already held and adds the provided
 208 // item to the queue if it does not already exist.
 209 func (h *Heap) addIfNotPresentLocked(key string, obj interface{}) {
 210         if _, exists := h.data.items[key]; exists {
 211                 return
 212         }
 213         heap.Push(h.data, &itemKeyValue{key, obj})
 214 }
 215
 216 // Update is the same as Add in this implementation. When the item does not
 217 // exist, it is added.
 218 func (h *Heap) Update(obj interface{}) error {
 219         return h.Add(obj)
 220 }
 221
 222 // Delete removes an item.
 223 func (h *Heap) Delete(obj interface{}) error {
 224         key, err := h.data.keyFunc(obj)
 225         if err != nil {
 226                 return KeyError{obj, err}
 227         }
 228         h.lock.Lock()
 229         defer h.lock.Unlock()
 230         if item, ok := h.data.items[key]; ok {
 231                 heap.Remove(h.data, item.index)
 232                 return nil
 233         }
 234         return fmt.Errorf("object not found")
 235 }
 236
 237 // Pop waits until an item is ready. If multiple items are
 238 // ready, they are returned in the order given by Heap.data.lessFunc.
 239 func (h *Heap) Pop() (interface{}, error) {
 240         h.lock.Lock()
 241         defer h.lock.Unlock()
 242         for len(h.data.queue) == 0 {
 243                 // When the queue is empty, invocation of Pop() is blocked until new item is enqueued.
 244                 // When Close() is called, the h.closed is set and the condition is broadcast,
 245                 // which causes this loop to continue and return from the Pop().
 246                 if h.closed {
 247                         return nil, fmt.Errorf("heap is closed")
 248                 }
 249                 h.cond.Wait()
 250         }
 251         obj := heap.Pop(h.data)
 252         if obj != nil {
 253                 return obj, nil
 254         } else {
 255                 return nil, fmt.Errorf("object was removed from heap data")
 256         }
 257 }
 258
 259 // List returns a list of all the items.
 260 func (h *Heap) List() []interface{} {
 261         h.lock.RLock()
 262         defer h.lock.RUnlock()
 263         list := make([]interface{}, 0, len(h.data.items))
 264         for _, item := range h.data.items {
 265                 list = append(list, item.obj)
 266         }
 267         return list
 268 }
 269
 270 // ListKeys returns a list of all the keys of the objects currently in the Heap.
 271 func (h *Heap) ListKeys() []string {
 272         h.lock.RLock()
 273         defer h.lock.RUnlock()
 274         list := make([]string, 0, len(h.data.items))
 275         for key := range h.data.items {
 276                 list = append(list, key)
 277         }
 278         return list
 279 }
 280
 281 // Get returns the requested item, or sets exists=false.
 282 func (h *Heap) Get(obj interface{}) (interface{}, bool, error) {
 283         key, err := h.data.keyFunc(obj)
 284         if err != nil {
 285                 return nil, false, KeyError{obj, err}
 286         }
 287         return h.GetByKey(key)
 288 }
 289
 290 // GetByKey returns the requested item, or sets exists=false.
 291 func (h *Heap) GetByKey(key string) (interface{}, bool, error) {
 292         h.lock.RLock()
 293         defer h.lock.RUnlock()
 294         item, exists := h.data.items[key]
 295         if !exists {
 296                 return nil, false, nil
 297         }
 298         return item.obj, true, nil
 299 }
 300
 301 // IsClosed returns true if the queue is closed.
 302 func (h *Heap) IsClosed() bool {
 303         h.lock.RLock()
 304         defer h.lock.RUnlock()
 305         if h.closed {
 306                 return true
 307         }
 308         return false
 309 }
 310
 311 // NewHeap returns a Heap which can be used to queue up items to process.
 312 func NewHeap(keyFn KeyFunc, lessFn LessFunc) *Heap {
 313         h := &Heap{
 314                 data: &heapData{
 315                         items:    map[string]*heapItem{},
 316                         queue:    []string{},
 317                         keyFunc:  keyFn,
 318                         lessFunc: lessFn,
 319                 },
 320         }
 321         h.cond.L = &h.lock
 322         return h
 323 }