diff --git a/tools/cache/BUILD b/tools/cache/BUILD index 684d3dca..75d4ff07 100644 --- a/tools/cache/BUILD +++ b/tools/cache/BUILD @@ -15,6 +15,7 @@ go_test( "delta_fifo_test.go", "expiration_cache_test.go", "fifo_test.go", + "heap_test.go", "index_test.go", "mutation_detector_test.go", "processor_listener_test.go", @@ -50,6 +51,7 @@ go_library( "expiration_cache_fakes.go", "fake_custom_store.go", "fifo.go", + "heap.go", "index.go", "listers.go", "listwatch.go", diff --git a/tools/cache/heap.go b/tools/cache/heap.go new file mode 100644 index 00000000..78e49245 --- /dev/null +++ b/tools/cache/heap.go @@ -0,0 +1,323 @@ +/* +Copyright 2017 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// This file implements a heap data structure. + +package cache + +import ( + "container/heap" + "fmt" + "sync" +) + +const ( + closedMsg = "heap is closed" +) + +type LessFunc func(interface{}, interface{}) bool +type heapItem struct { + obj interface{} // The object which is stored in the heap. + index int // The index of the object's key in the Heap.queue. +} + +type itemKeyValue struct { + key string + obj interface{} +} + +// heapData is an internal struct that implements the standard heap interface +// and keeps the data stored in the heap. +type heapData struct { + // items is a map from key of the objects to the objects and their index. + // We depend on the property that items in the map are in the queue and vice versa. + items map[string]*heapItem + // queue implements a heap data structure and keeps the order of elements + // according to the heap invariant. The queue keeps the keys of objects stored + // in "items". + queue []string + + // keyFunc is used to make the key used for queued item insertion and retrieval, and + // should be deterministic. + keyFunc KeyFunc + // lessFunc is used to compare two objects in the heap. + lessFunc LessFunc +} + +var ( + _ = heap.Interface(&heapData{}) // heapData is a standard heap +) + +// Less compares two objects and returns true if the first one should go +// in front of the second one in the heap. +func (h *heapData) Less(i, j int) bool { + if i > len(h.queue) || j > len(h.queue) { + return false + } + itemi, ok := h.items[h.queue[i]] + if !ok { + return false + } + itemj, ok := h.items[h.queue[j]] + if !ok { + return false + } + return h.lessFunc(itemi.obj, itemj.obj) +} + +// Len returns the number of items in the Heap. +func (h *heapData) Len() int { return len(h.queue) } + +// Swap implements swapping of two elements in the heap. This is a part of standard +// heap interface and should never be called directly. +func (h *heapData) Swap(i, j int) { + h.queue[i], h.queue[j] = h.queue[j], h.queue[i] + item := h.items[h.queue[i]] + item.index = i + item = h.items[h.queue[j]] + item.index = j +} + +// Push is supposed to be called by heap.Push only. +func (h *heapData) Push(kv interface{}) { + keyValue := kv.(*itemKeyValue) + n := len(h.queue) + h.items[keyValue.key] = &heapItem{keyValue.obj, n} + h.queue = append(h.queue, keyValue.key) +} + +// Pop is supposed to be called by heap.Pop only. +func (h *heapData) Pop() interface{} { + key := h.queue[len(h.queue)-1] + h.queue = h.queue[0 : len(h.queue)-1] + item, ok := h.items[key] + if !ok { + // This is an error + return nil + } + delete(h.items, key) + return item.obj +} + +// Heap is a thread-safe producer/consumer queue that implements a heap data structure. +// It can be used to implement priority queues and similar data structures. +type Heap struct { + lock sync.RWMutex + cond sync.Cond + + // data stores objects and has a queue that keeps their ordering according + // to the heap invariant. + data *heapData + + // closed indicates that the queue is closed. + // It is mainly used to let Pop() exit its control loop while waiting for an item. + closed bool +} + +// Close the Heap and signals condition variables that may be waiting to pop +// items from the heap. +func (h *Heap) Close() { + h.lock.Lock() + defer h.lock.Unlock() + h.closed = true + h.cond.Broadcast() +} + +// Add inserts an item, and puts it in the queue. The item is updated if it +// already exists. +func (h *Heap) Add(obj interface{}) error { + key, err := h.data.keyFunc(obj) + if err != nil { + return KeyError{obj, err} + } + h.lock.Lock() + defer h.lock.Unlock() + if h.closed { + return fmt.Errorf(closedMsg) + } + if _, exists := h.data.items[key]; exists { + h.data.items[key].obj = obj + heap.Fix(h.data, h.data.items[key].index) + } else { + h.addIfNotPresentLocked(key, obj) + } + h.cond.Broadcast() + return nil +} + +// Adds all the items in the list to the queue and then signals the condition +// variable. It is useful when the caller would like to add all of the items +// to the queue before consumer starts processing them. +func (h *Heap) BulkAdd(list []interface{}) error { + h.lock.Lock() + defer h.lock.Unlock() + if h.closed { + return fmt.Errorf(closedMsg) + } + for _, obj := range list { + key, err := h.data.keyFunc(obj) + if err != nil { + return KeyError{obj, err} + } + if _, exists := h.data.items[key]; exists { + h.data.items[key].obj = obj + heap.Fix(h.data, h.data.items[key].index) + } else { + h.addIfNotPresentLocked(key, obj) + } + } + h.cond.Broadcast() + return nil +} + +// AddIfNotPresent inserts an item, and puts it in the queue. If an item with +// the key is present in the map, no changes is made to the item. +// +// This is useful in a single producer/consumer scenario so that the consumer can +// safely retry items without contending with the producer and potentially enqueueing +// stale items. +func (h *Heap) AddIfNotPresent(obj interface{}) error { + id, err := h.data.keyFunc(obj) + if err != nil { + return KeyError{obj, err} + } + h.lock.Lock() + defer h.lock.Unlock() + if h.closed { + return fmt.Errorf(closedMsg) + } + h.addIfNotPresentLocked(id, obj) + h.cond.Broadcast() + return nil +} + +// addIfNotPresentLocked assumes the lock is already held and adds the the provided +// item to the queue if it does not already exist. +func (h *Heap) addIfNotPresentLocked(key string, obj interface{}) { + if _, exists := h.data.items[key]; exists { + return + } + heap.Push(h.data, &itemKeyValue{key, obj}) +} + +// Update is the same as Add in this implementation. When the item does not +// exist, it is added. +func (h *Heap) Update(obj interface{}) error { + return h.Add(obj) +} + +// Delete removes an item. +func (h *Heap) Delete(obj interface{}) error { + key, err := h.data.keyFunc(obj) + if err != nil { + return KeyError{obj, err} + } + h.lock.Lock() + defer h.lock.Unlock() + if item, ok := h.data.items[key]; ok { + heap.Remove(h.data, item.index) + return nil + } + return fmt.Errorf("object not found") +} + +// Pop waits until an item is ready. If multiple items are +// ready, they are returned in the order given by Heap.data.lessFunc. +func (h *Heap) Pop() (interface{}, error) { + h.lock.Lock() + defer h.lock.Unlock() + for len(h.data.queue) == 0 { + // When the queue is empty, invocation of Pop() is blocked until new item is enqueued. + // When Close() is called, the h.closed is set and the condition is broadcast, + // which causes this loop to continue and return from the Pop(). + if h.closed { + return nil, fmt.Errorf("heap is closed") + } + h.cond.Wait() + } + obj := heap.Pop(h.data) + if obj != nil { + return obj, nil + } else { + return nil, fmt.Errorf("object was removed from heap data") + } +} + +// List returns a list of all the items. +func (h *Heap) List() []interface{} { + h.lock.RLock() + defer h.lock.RUnlock() + list := make([]interface{}, 0, len(h.data.items)) + for _, item := range h.data.items { + list = append(list, item.obj) + } + return list +} + +// ListKeys returns a list of all the keys of the objects currently in the Heap. +func (h *Heap) ListKeys() []string { + h.lock.RLock() + defer h.lock.RUnlock() + list := make([]string, 0, len(h.data.items)) + for key := range h.data.items { + list = append(list, key) + } + return list +} + +// Get returns the requested item, or sets exists=false. +func (h *Heap) Get(obj interface{}) (interface{}, bool, error) { + key, err := h.data.keyFunc(obj) + if err != nil { + return nil, false, KeyError{obj, err} + } + return h.GetByKey(key) +} + +// GetByKey returns the requested item, or sets exists=false. +func (h *Heap) GetByKey(key string) (interface{}, bool, error) { + h.lock.RLock() + defer h.lock.RUnlock() + item, exists := h.data.items[key] + if !exists { + return nil, false, nil + } + return item.obj, true, nil +} + +// IsClosed returns true if the queue is closed. +func (h *Heap) IsClosed() bool { + h.lock.RLock() + defer h.lock.RUnlock() + if h.closed { + return true + } + return false +} + +// NewHeap returns a Heap which can be used to queue up items to process. +func NewHeap(keyFn KeyFunc, lessFn LessFunc) *Heap { + h := &Heap{ + data: &heapData{ + items: map[string]*heapItem{}, + queue: []string{}, + keyFunc: keyFn, + lessFunc: lessFn, + }, + } + h.cond.L = &h.lock + return h +} diff --git a/tools/cache/heap_test.go b/tools/cache/heap_test.go new file mode 100644 index 00000000..c2e47698 --- /dev/null +++ b/tools/cache/heap_test.go @@ -0,0 +1,382 @@ +/* +Copyright 2017 The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package cache + +import ( + "sync" + "testing" + "time" +) + +func testHeapObjectKeyFunc(obj interface{}) (string, error) { + return obj.(testHeapObject).name, nil +} + +type testHeapObject struct { + name string + val interface{} +} + +func mkHeapObj(name string, val interface{}) testHeapObject { + return testHeapObject{name: name, val: val} +} + +func compareInts(val1 interface{}, val2 interface{}) bool { + first := val1.(testHeapObject).val.(int) + second := val2.(testHeapObject).val.(int) + return first < second +} + +// TestHeapBasic tests Heap invariant and synchronization. +func TestHeapBasic(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + var wg sync.WaitGroup + wg.Add(2) + const amount = 500 + var i, u int + // Insert items in the heap in opposite orders in two go routines. + go func() { + for i = amount; i > 0; i-- { + h.Add(mkHeapObj(string([]rune{'a', rune(i)}), i)) + } + wg.Done() + }() + go func() { + for u = 0; u < amount; u++ { + h.Add(mkHeapObj(string([]rune{'b', rune(u)}), u+1)) + } + wg.Done() + }() + // Wait for the two go routines to finish. + wg.Wait() + // Make sure that the numbers are popped in ascending order. + prevNum := 0 + for i := 0; i < amount*2; i++ { + obj, err := h.Pop() + num := obj.(testHeapObject).val.(int) + // All the items must be sorted. + if err != nil || prevNum > num { + t.Errorf("got %v out of order, last was %v", obj, prevNum) + } + prevNum = num + } +} + +// Tests Heap.Add and ensures that heap invariant is preserved after adding items. +func TestHeap_Add(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + h.Add(mkHeapObj("baz", 11)) + h.Add(mkHeapObj("zab", 30)) + h.Add(mkHeapObj("foo", 13)) // This updates "foo". + + item, err := h.Pop() + if e, a := 1, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + item, err = h.Pop() + if e, a := 11, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + h.Delete(mkHeapObj("baz", 11)) // Nothing is deleted. + h.Add(mkHeapObj("foo", 14)) // foo is updated. + item, err = h.Pop() + if e, a := 14, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + item, err = h.Pop() + if e, a := 30, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } +} + +// TestHeap_BulkAdd tests Heap.BulkAdd functionality and ensures that all the +// items given to BulkAdd are added to the queue before Pop reads them. +func TestHeap_BulkAdd(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + const amount = 500 + // Insert items in the heap in opposite orders in a go routine. + go func() { + l := []interface{}{} + for i := amount; i > 0; i-- { + l = append(l, mkHeapObj(string([]rune{'a', rune(i)}), i)) + } + h.BulkAdd(l) + }() + prevNum := -1 + for i := 0; i < amount; i++ { + obj, err := h.Pop() + num := obj.(testHeapObject).val.(int) + // All the items must be sorted. + if err != nil || prevNum >= num { + t.Errorf("got %v out of order, last was %v", obj, prevNum) + } + prevNum = num + } +} + +// TestHeapEmptyPop tests that pop returns properly after heap is closed. +func TestHeapEmptyPop(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + go func() { + time.Sleep(1 * time.Second) + h.Close() + }() + _, err := h.Pop() + if err == nil || err.Error() != closedMsg { + t.Errorf("pop should have returned heap closed error: %v", err) + } +} + +// TestHeap_AddIfNotPresent tests Heap.AddIfNotPresent and ensures that heap +// invariant is preserved after adding items. +func TestHeap_AddIfNotPresent(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.AddIfNotPresent(mkHeapObj("foo", 10)) + h.AddIfNotPresent(mkHeapObj("bar", 1)) + h.AddIfNotPresent(mkHeapObj("baz", 11)) + h.AddIfNotPresent(mkHeapObj("zab", 30)) + h.AddIfNotPresent(mkHeapObj("foo", 13)) // This is not added. + + if len := len(h.data.items); len != 4 { + t.Errorf("unexpected number of items: %d", len) + } + if val := h.data.items["foo"].obj.(testHeapObject).val; val != 10 { + t.Errorf("unexpected value: %d", val) + } + item, err := h.Pop() + if e, a := 1, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + item, err = h.Pop() + if e, a := 10, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + // bar is already popped. Let's add another one. + h.AddIfNotPresent(mkHeapObj("bar", 14)) + item, err = h.Pop() + if e, a := 11, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + item, err = h.Pop() + if e, a := 14, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } +} + +// TestHeap_Delete tests Heap.Delete and ensures that heap invariant is +// preserved after deleting items. +func TestHeap_Delete(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + h.Add(mkHeapObj("bal", 31)) + h.Add(mkHeapObj("baz", 11)) + + // Delete head. Delete should work with "key" and doesn't care about the value. + if err := h.Delete(mkHeapObj("bar", 200)); err != nil { + t.Fatalf("Failed to delete head.") + } + item, err := h.Pop() + if e, a := 10, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + h.Add(mkHeapObj("zab", 30)) + h.Add(mkHeapObj("faz", 30)) + len := h.data.Len() + // Delete non-existing item. + if err = h.Delete(mkHeapObj("non-existent", 10)); err == nil || len != h.data.Len() { + t.Fatalf("Didn't expect any item removal") + } + // Delete tail. + if err = h.Delete(mkHeapObj("bal", 31)); err != nil { + t.Fatalf("Failed to delete tail.") + } + // Delete one of the items with value 30. + if err = h.Delete(mkHeapObj("zab", 30)); err != nil { + t.Fatalf("Failed to delete item.") + } + item, err = h.Pop() + if e, a := 11, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + item, err = h.Pop() + if e, a := 30, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + if h.data.Len() != 0 { + t.Fatalf("expected an empty heap.") + } +} + +// TestHeap_Update tests Heap.Update and ensures that heap invariant is +// preserved after adding items. +func TestHeap_Update(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + h.Add(mkHeapObj("bal", 31)) + h.Add(mkHeapObj("baz", 11)) + + // Update an item to a value that should push it to the head. + h.Update(mkHeapObj("baz", 0)) + if h.data.queue[0] != "baz" || h.data.items["baz"].index != 0 { + t.Fatalf("expected baz to be at the head") + } + item, err := h.Pop() + if e, a := 0, item.(testHeapObject).val; err != nil || a != e { + t.Fatalf("expected %d, got %d", e, a) + } + // Update bar to push it farther back in the queue. + h.Update(mkHeapObj("bar", 100)) + if h.data.queue[0] != "foo" || h.data.items["foo"].index != 0 { + t.Fatalf("expected foo to be at the head") + } +} + +// TestHeap_Get tests Heap.Get. +func TestHeap_Get(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + h.Add(mkHeapObj("bal", 31)) + h.Add(mkHeapObj("baz", 11)) + + // Get works with the key. + obj, exists, err := h.Get(mkHeapObj("baz", 0)) + if err != nil || exists == false || obj.(testHeapObject).val != 11 { + t.Fatalf("unexpected error in getting element") + } + // Get non-existing object. + _, exists, err = h.Get(mkHeapObj("non-existing", 0)) + if err != nil || exists == true { + t.Fatalf("didn't expect to get any object") + } +} + +// TestHeap_GetByKey tests Heap.GetByKey and is very similar to TestHeap_Get. +func TestHeap_GetByKey(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + h.Add(mkHeapObj("bal", 31)) + h.Add(mkHeapObj("baz", 11)) + + obj, exists, err := h.GetByKey("baz") + if err != nil || exists == false || obj.(testHeapObject).val != 11 { + t.Fatalf("unexpected error in getting element") + } + // Get non-existing object. + _, exists, err = h.GetByKey("non-existing") + if err != nil || exists == true { + t.Fatalf("didn't expect to get any object") + } +} + +// TestHeap_Close tests Heap.Close and Heap.IsClosed functions. +func TestHeap_Close(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Add(mkHeapObj("foo", 10)) + h.Add(mkHeapObj("bar", 1)) + + if h.IsClosed() { + t.Fatalf("didn't expect heap to be closed") + } + h.Close() + if !h.IsClosed() { + t.Fatalf("expect heap to be closed") + } +} + +// TestHeap_List tests Heap.List function. +func TestHeap_List(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + list := h.List() + if len(list) != 0 { + t.Errorf("expected an empty list") + } + + items := map[string]int{ + "foo": 10, + "bar": 1, + "bal": 30, + "baz": 11, + "faz": 30, + } + for k, v := range items { + h.Add(mkHeapObj(k, v)) + } + list = h.List() + if len(list) != len(items) { + t.Errorf("expected %d items, got %d", len(items), len(list)) + } + for _, obj := range list { + heapObj := obj.(testHeapObject) + v, ok := items[heapObj.name] + if !ok || v != heapObj.val { + t.Errorf("unexpected item in the list: %v", heapObj) + } + } +} + +// TestHeap_ListKeys tests Heap.ListKeys function. Scenario is the same as +// TestHeap_list. +func TestHeap_ListKeys(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + list := h.ListKeys() + if len(list) != 0 { + t.Errorf("expected an empty list") + } + + items := map[string]int{ + "foo": 10, + "bar": 1, + "bal": 30, + "baz": 11, + "faz": 30, + } + for k, v := range items { + h.Add(mkHeapObj(k, v)) + } + list = h.ListKeys() + if len(list) != len(items) { + t.Errorf("expected %d items, got %d", len(items), len(list)) + } + for _, key := range list { + _, ok := items[key] + if !ok { + t.Errorf("unexpected item in the list: %v", key) + } + } +} + +// TestHeapAddAfterClose tests that heap returns an error if anything is added +// after it is closed. +func TestHeapAddAfterClose(t *testing.T) { + h := NewHeap(testHeapObjectKeyFunc, compareInts) + h.Close() + if err := h.Add(mkHeapObj("test", 1)); err == nil || err.Error() != closedMsg { + t.Errorf("expected heap closed error") + } + if err := h.AddIfNotPresent(mkHeapObj("test", 1)); err == nil || err.Error() != closedMsg { + t.Errorf("expected heap closed error") + } + if err := h.BulkAdd([]interface{}{mkHeapObj("test", 1)}); err == nil || err.Error() != closedMsg { + t.Errorf("expected heap closed error") + } +}