fix

pkg/pqueue/pqueue.go

@@ -0,0 +1,89 @@
+// Package pqueue jiacrontab中使用的优先队列
+// 参考nsq的实现
+// 做了注释和少量调整
+package pqueue
+
+import (
+	"container/heap"
+)
+
+type Item struct {
+	Value    interface{}
+	Priority int64
+	Index    int
+}
+
+// PriorityQueue 最小堆实现的优先队列
+type PriorityQueue []*Item
+
+// New 创建
+func New(capacity int) PriorityQueue {
+	return make(PriorityQueue, 0, capacity)
+}
+
+// Len 队列长队
+func (pq PriorityQueue) Len() int {
+	return len(pq)
+}
+
+// Less 比较相邻两个原素优先级
+func (pq PriorityQueue) Less(i, j int) bool {
+	return pq[i].Priority < pq[j].Priority
+}
+
+// Swap 交换相邻原素
+func (pq PriorityQueue) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].Index = i
+	pq[j].Index = j
+}
+
+// Push 添加新的item
+func (pq *PriorityQueue) Push(x interface{}) {
+	n := len(*pq)
+	c := cap(*pq)
+	if n+1 > c {
+		npq := make(PriorityQueue, n, c*2)
+		copy(npq, *pq)
+		*pq = npq
+	}
+	*pq = (*pq)[0 : n+1]
+	item := x.(*Item)
+	item.Index = n
+	(*pq)[n] = item
+}
+
+func (pq *PriorityQueue) update(item *Item, value string, priority int64) {
+	item.Value = value
+	item.Priority = priority
+	heap.Fix(pq, item.Index)
+}
+
+// Pop 弹出队列末端原素
+func (pq *PriorityQueue) Pop() interface{} {
+	n := len(*pq)
+	c := cap(*pq)
+	if n < (c/2) && c > 25 {
+		npq := make(PriorityQueue, n, c/2)
+		copy(npq, *pq)
+		*pq = npq
+	}
+	item := (*pq)[n-1]
+	item.Index = -1
+	*pq = (*pq)[0 : n-1]
+	return item
+}
+
+// PeekAndShift 根据比较max并弹出原素
+func (pq *PriorityQueue) PeekAndShift(max int64) (*Item, int64) {
+	if pq.Len() == 0 {
+		return nil, 0
+	}
+
+	item := (*pq)[0]
+	if item.Priority > max {
+		return nil, item.Priority - max
+	}
+	heap.Remove(pq, 0)
+	return item, 0
+}

pkg/pqueue/pqueue_test.go

@@ -0,0 +1,79 @@
+package pqueue
+
+import (
+	"container/heap"
+	"math/rand"
+	"path/filepath"
+	"reflect"
+	"runtime"
+	"sort"
+	"testing"
+)
+
+func equal(t *testing.T, act, exp interface{}) {
+	if !reflect.DeepEqual(exp, act) {
+		_, file, line, _ := runtime.Caller(1)
+		t.Logf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n",
+			filepath.Base(file), line, exp, act)
+		t.FailNow()
+	}
+}
+
+func TestPriorityQueue(t *testing.T) {
+	c := 100
+	pq := New(c)
+
+	for i := 0; i < c+1; i++ {
+		heap.Push(&pq, &Item{Value: i, Priority: int64(i)})
+	}
+	equal(t, pq.Len(), c+1)
+	equal(t, cap(pq), c*2)
+
+	for i := 0; i < c+1; i++ {
+		item := heap.Pop(&pq)
+		equal(t, item.(*Item).Value.(int), i)
+	}
+	equal(t, cap(pq), c/4)
+}
+
+func TestUnsortedInsert(t *testing.T) {
+	c := 100
+	pq := New(c)
+	ints := make([]int, 0, c)
+
+	for i := 0; i < c; i++ {
+		v := rand.Int()
+		ints = append(ints, v)
+		heap.Push(&pq, &Item{Value: i, Priority: int64(v)})
+	}
+	equal(t, pq.Len(), c)
+	equal(t, cap(pq), c)
+
+	sort.Sort(sort.IntSlice(ints))
+
+	for i := 0; i < c; i++ {
+		item, _ := pq.PeekAndShift(int64(ints[len(ints)-1]))
+		equal(t, item.Priority, int64(ints[i]))
+	}
+}
+
+func TestRemove(t *testing.T) {
+	c := 100
+	pq := New(c)
+
+	for i := 0; i < c; i++ {
+		v := rand.Int()
+		heap.Push(&pq, &Item{Value: "test", Priority: int64(v)})
+	}
+
+	for i := 0; i < 10; i++ {
+		heap.Remove(&pq, rand.Intn((c-1)-i))
+	}
+
+	lastPriority := heap.Pop(&pq).(*Item).Priority
+	for i := 0; i < (c - 10 - 1); i++ {
+		item := heap.Pop(&pq)
+		equal(t, lastPriority < item.(*Item).Priority, true)
+		lastPriority = item.(*Item).Priority
+	}
+}