diff --git a/queue/priority_queue.go b/queue/priority_queue.go new file mode 100644 index 000000000..aae7b423c --- /dev/null +++ b/queue/priority_queue.go @@ -0,0 +1,74 @@ +package queue + +import ( + "container/heap" +) + +// PriorityQueueItem is an interface that represents items in a PriorityQueue. +// Users of PriorityQueue will need to define a Less function such that +// PriorityQueue will be able to use that to build and restore an underlying +// heap. +type PriorityQueueItem interface { + Less(other PriorityQueueItem) bool +} + +type priorityQueue []PriorityQueueItem + +// Len returns the length of the priorityQueue. +func (pq priorityQueue) Len() int { return len(pq) } + +// Less is used to order PriorityQueueItem items in the queue. +func (pq priorityQueue) Less(i, j int) bool { + return pq[i].Less(pq[j]) +} + +// Swap swaps two items in the priorityQueue. Swap is used by heap.Interface. +func (pq priorityQueue) Swap(i, j int) { + pq[i], pq[j] = pq[j], pq[i] +} + +// Push adds a new item the the priorityQueue. +func (pq *priorityQueue) Push(x interface{}) { + item := x.(PriorityQueueItem) + *pq = append(*pq, item) +} + +// Pop removes the top item from the priorityQueue. +func (pq *priorityQueue) Pop() interface{} { + old := *pq + n := len(old) + item := old[n-1] + old[n-1] = nil + *pq = old[0 : n-1] + return item +} + +// Priority wrap a standard heap in a more object-oriented structure. +type PriorityQueue struct { + queue priorityQueue +} + +// Len returns the length of the queue. +func (pq *PriorityQueue) Len() int { + return len(pq.queue) +} + +// Empty returns true if the queue is empty. +func (pq *PriorityQueue) Empty() bool { + return len(pq.queue) == 0 +} + +// Push adds an item to the priority queue. +func (pq *PriorityQueue) Push(item PriorityQueueItem) { + heap.Push(&pq.queue, item) +} + +// Pop removes the top most item from the queue. +func (pq *PriorityQueue) Pop() PriorityQueueItem { + return heap.Pop(&pq.queue).(PriorityQueueItem) +} + +// Top returns the top most item from the queue without removing it. +func (pq *PriorityQueue) Top() PriorityQueueItem { + return pq.queue[0] +} diff --git a/queue/priority_queue_test.go b/queue/priority_queue_test.go new file mode 100644 index 000000000..be1209e06 --- /dev/null +++ b/queue/priority_queue_test.go @@ -0,0 +1,67 @@ +package queue + +import ( + "math/rand" + "testing" + "time" +) + +type testQueueItem struct { + Value int + Expiry time.Time +} + +func (e testQueueItem) Less(other PriorityQueueItem) bool { + return e.Expiry.Before(other.(*testQueueItem).Expiry) +} + +func TestExpiryQueue(t *testing.T) { + // The number of elements we push to the queue. + count := 100 + // Generate a random permutation of a range [0, count) + array := rand.Perm(count) + // t0 holds a reference time point. + t0 := time.Date(1975, time.April, 5, 12, 0, 0, 0, time.UTC) + + var testQueue PriorityQueue + + if testQueue.Len() != 0 && !testQueue.Empty() { + t.Fatal("Expected the queue to be empty") + } + + // Create elements with expiry of t0 + value * second. + for _, value := range array { + testQueue.Push(&testQueueItem{ + Value: value, + Expiry: t0.Add(time.Duration(value) * time.Second), + }) + } + + // Now expect that we can retrieve elements in order of their expiry. + for i := 0; i < count; i++ { + expectedQueueLen := count - i + if testQueue.Len() != expectedQueueLen { + t.Fatalf("Expected the queue len %v, got %v", + expectedQueueLen, testQueue.Len()) + } + + if testQueue.Empty() { + t.Fatalf("Did not expect the queue to be empty") + } + + top := testQueue.Top().(*testQueueItem) + if top.Value != i { + t.Fatalf("Expected queue top %v, got %v", i, top.Value) + } + + popped := testQueue.Pop().(*testQueueItem) + if popped != top { + t.Fatalf("Expected queue top %v equal to popped: %v", + top, popped) + } + } + + if testQueue.Len() != 0 || !testQueue.Empty() { + t.Fatalf("Expected the queue to be empty") + } +}