Merge pull request #49752 from bsalamat/priority_scheduler

Automatic merge from submit-queue

Add a heap to client-go. Heap orders items with heap invariant ordering.

**What this PR does / why we need it**:
Heap is useful in implementing priority queues. Some components may need such ordering to process their highest priority objects first. Scheduler is going to be the first user of the heap. It will store pending pods ordered by their priority, so that the highest priority pods are popped first to be scheduled.

**Which issue this PR fixes** *(optional, in `fixes #<issue number>(, fixes #<issue_number>, ...)` format, will close that issue when PR gets merged)*: fixes #

**Special notes for your reviewer**:

**Release note**:

```release-note
NONE
```

ref/ #47604
ref/ #48646

@kubernetes/api-reviewers @kubernetes/sig-scheduling-pr-reviews @davidopp 
/assign @caesarxuchao

Kubernetes-commit: 868fef189c8de9fcb61039f064fe0a02a3b06198

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",

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
+}

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")
+	}
+}