Browse Source

feat add ttl heap

Xiang Li 12 years ago
parent
commit
5394c8a894
2 changed files with 85 additions and 0 deletions
  1. 36 0
      store/heap_test.go
  2. 49 0
      store/ttl_key_heap.go

+ 36 - 0
store/heap_test.go

@@ -0,0 +1,36 @@
+package store
+
+import (
+	"container/heap"
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestHeapPushPop(t *testing.T) {
+	h := &TTLKeyHeap{Map: make(map[*Node]int)}
+	heap.Init(h)
+
+	kvs := make([]*Node, 10)
+
+	// add from older expire time to earlier expire time
+	// the path is equal to ttl from now
+	for i, n := range kvs {
+		path := fmt.Sprintf("%v", 10-i)
+		m := time.Duration(10 - i)
+		n = newKV(nil, path, path, 0, 0, nil, "", time.Now().Add(time.Second*m))
+		heap.Push(h, n)
+	}
+
+	min := time.Now()
+
+	for i := 0; i < 9; i++ {
+		iNode := heap.Pop(h)
+		node, _ := iNode.(*Node)
+		if node.ExpireTime.Before(min) {
+			t.Fatal("heap sort wrong!")
+		}
+		min = node.ExpireTime
+	}
+
+}

+ 49 - 0
store/ttl_key_heap.go

@@ -0,0 +1,49 @@
+package store
+
+import (
+	"container/heap"
+)
+
+// An TTLKeyHeap is a min-heap of TTLKeys order by expiration time
+type TTLKeyHeap struct {
+	Array []*Node
+	Map   map[*Node]int
+}
+
+func (h TTLKeyHeap) Len() int {
+	return len(h.Array)
+}
+
+func (h TTLKeyHeap) Less(i, j int) bool {
+	return h.Array[i].ExpireTime.Before(h.Array[j].ExpireTime)
+}
+
+func (h TTLKeyHeap) Swap(i, j int) {
+	// swap node
+	h.Array[i], h.Array[j] = h.Array[j], h.Array[i]
+
+	// update map
+	h.Map[h.Array[i]] = i
+	h.Map[h.Array[j]] = j
+}
+
+func (h *TTLKeyHeap) Push(x interface{}) {
+	n, _ := x.(*Node)
+	h.Map[n] = len(h.Array)
+	h.Array = append(h.Array, n)
+}
+
+func (h *TTLKeyHeap) Pop() interface{} {
+	old := h.Array
+	n := len(old)
+	x := old[n-1]
+	h.Array = old[0 : n-1]
+	delete(h.Map, x)
+	return x
+}
+
+func (h *TTLKeyHeap) Update(n *Node) {
+	index := h.Map[n]
+	heap.Remove(h, index)
+	heap.Push(h, n)
+}