diff --git a/store/heap_test.go b/store/heap_test.go new file mode 100644 index 000000000..02682b973 --- /dev/null +++ b/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 + } + +} diff --git a/store/ttl_key_heap.go b/store/ttl_key_heap.go new file mode 100644 index 000000000..abb12fede --- /dev/null +++ b/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) +}